From 111994956ede6964433e02bc3f64deb04b9916ac Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 14 Jun 2020 19:25:39 -0700
Subject: [PATCH] switch default inference to FP16 on GPU

---
 detect.py |  7 +++----
 test.py   | 22 ++++++++++------------
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/detect.py b/detect.py
index fa67269..86b6169 100644
--- a/detect.py
+++ b/detect.py
@@ -5,8 +5,8 @@ from utils.utils import *
 
 
 def detect(save_img=False):
-    out, source, weights, half, view_img, save_txt, imgsz = \
-        opt.output, opt.source, opt.weights, opt.half, opt.view_img, opt.save_txt, opt.img_size
+    out, source, weights, view_img, save_txt, imgsz = \
+        opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
     webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')
 
     # Initialize
@@ -14,7 +14,7 @@ def detect(save_img=False):
     if os.path.exists(out):
         shutil.rmtree(out)  # delete output folder
     os.makedirs(out)  # make new output folder
-    half &= device.type != 'cpu'  # half precision only supported on CUDA
+    half = device.type != 'cpu'  # half precision only supported on CUDA
 
     # Load model
     google_utils.attempt_download(weights)
@@ -142,7 +142,6 @@ if __name__ == '__main__':
     parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')
     parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
     parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
-    parser.add_argument('--half', action='store_true', help='half precision FP16 inference')
     parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
     parser.add_argument('--view-img', action='store_true', help='display results')
     parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
diff --git a/test.py b/test.py
index f5bf3df..2f031ab 100644
--- a/test.py
+++ b/test.py
@@ -17,7 +17,6 @@ def test(data,
          save_json=False,
          single_cls=False,
          augment=False,
-         half=False,  # FP16
          model=None,
          dataloader=None,
          fast=False,
@@ -25,7 +24,7 @@ def test(data,
     # Initialize/load model and set device
     if model is None:
         device = torch_utils.select_device(opt.device, batch_size=batch_size)
-        half &= device.type != 'cpu'  # half precision only supported on CUDA
+        half = device.type != 'cpu'  # half precision only supported on CUDA
 
         # Remove previous
         for f in glob.glob('test_batch*.jpg'):
@@ -48,7 +47,8 @@ def test(data,
         device = next(model.parameters()).device  # get model device
         training = True
 
-    # Configure run
+    # Configure
+    model.eval()
     with open(data) as f:
         data = yaml.load(f, Loader=yaml.FullLoader)  # model dict
     nc = 1 if single_cls else int(data['nc'])  # number of classes
@@ -57,7 +57,10 @@ def test(data,
     niou = iouv.numel()
 
     # Dataloader
-    if dataloader is None:
+    if dataloader is None:  # not training
+        img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
+        _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once
+
         fast |= conf_thres > 0.001  # enable fast mode
         path = data['test'] if opt.task == 'test' else data['val']  # path to val/test images
         dataset = LoadImagesAndLabels(path,
@@ -75,9 +78,6 @@ def test(data,
                                 collate_fn=dataset.collate_fn)
 
     seen = 0
-    model.eval()
-    img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
-    _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once
     names = model.names if hasattr(model, 'names') else model.module.names
     coco91class = coco80_to_coco91_class()
     s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
@@ -221,11 +221,11 @@ def test(data,
             cocoDt = cocoGt.loadRes(f)  # initialize COCO pred api
 
             cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
-            cocoEval.params.imgIds = imgIds  # [:32]  # only evaluate these images
+            cocoEval.params.imgIds = imgIds  # image IDs to evaluate
             cocoEval.evaluate()
             cocoEval.accumulate()
             cocoEval.summarize()
-            map, map50 = cocoEval.stats[:2]  # update to pycocotools results (mAP@0.5:0.95, mAP@0.5)
+            map, map50 = cocoEval.stats[:2]  # update results (mAP@0.5:0.95, mAP@0.5)
         except:
             print('WARNING: pycocotools must be installed with numpy==1.17 to run correctly. '
                   'See https://github.com/cocodataset/cocoapi/issues/356')
@@ -248,7 +248,6 @@ if __name__ == '__main__':
     parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
     parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
     parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
-    parser.add_argument('--half', action='store_true', help='half precision FP16 inference')
     parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
     parser.add_argument('--augment', action='store_true', help='augmented inference')
     parser.add_argument('--verbose', action='store_true', help='report mAP by class')
@@ -268,8 +267,7 @@ if __name__ == '__main__':
              opt.iou_thres,
              opt.save_json,
              opt.single_cls,
-             opt.augment,
-             opt.half)
+             opt.augment)
 
     elif opt.task == 'study':  # run over a range of settings and save/plot
         for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: