diff --git a/detect.py b/detect.py index c6b76d981541..aba87687e666 100644 --- a/detect.py +++ b/detect.py @@ -28,7 +28,7 @@ def detect(opt): # Initialize set_logging() device = select_device(opt.device) - half = device.type != 'cpu' # half precision only supported on CUDA + half = opt.half and device.type != 'cpu' # half precision only supported on CUDA # Load model model = attempt_load(weights, map_location=device) # load FP32 model @@ -172,6 +172,7 @@ def detect(opt): parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)') parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels') parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences') + parser.add_argument('--half', type=bool, default=False, help='use FP16 half-precision inference') opt = parser.parse_args() print(opt) check_requirements(exclude=('tensorboard', 'pycocotools', 'thop')) diff --git a/test.py b/test.py index 0716c5d8b93c..113316ff0b8b 100644 --- a/test.py +++ b/test.py @@ -306,6 +306,7 @@ def test(data, parser.add_argument('--project', default='runs/test', help='save to project/name') parser.add_argument('--name', default='exp', help='save to project/name') parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') + parser.add_argument('--half', type=bool, default=False, help='use FP16 half-precision inference') opt = parser.parse_args() opt.save_json |= opt.data.endswith('coco.yaml') opt.data = check_file(opt.data) # check file @@ -326,6 +327,7 @@ def test(data, save_txt=opt.save_txt | opt.save_hybrid, save_hybrid=opt.save_hybrid, save_conf=opt.save_conf, + half_precision=opt.half, opt=opt )