maxmustermann/DeepLearningExamples
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

[MaskRCNN/PyT] Update AMP API for inference (#810)

Browse source
This commit is contained in:
Sharath TS 2021-01-13 18:15:05 -08:00 committed by GitHub
parent 2badf6e8e4
commit 0cbadd7d49
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
PyTorch/Segmentation/MaskRCNN/pytorch/tools/test_net.py
View file

@ -96,13 +96,14 @@ def main():
model = build_detection_model(cfg)
model.to(cfg.MODEL.DEVICE)
# Initialize mixed-precision if necessary
# Initialize mixed-precision
if args.fp16:
use_mixed_precision = True
else:
use_mixed_precision = cfg.DTYPE == "float16"
amp_handle = amp.init(enabled=use_mixed_precision, verbose=cfg.AMP_VERBOSE)
amp_opt_level = 'O1' if use_mixed_precision else 'O0'
model = amp.initialize(model, opt_level=amp_opt_level)
output_dir = cfg.OUTPUT_DIR
checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir)
_ = checkpointer.load(cfg.MODEL.WEIGHT)