glenn-jocher commited on
Commit
1119949
·
1 Parent(s): 157aff2

switch default inference to FP16 on GPU

Browse files
Files changed (2) hide show
  1. detect.py +3 -4
  2. test.py +10 -12
detect.py CHANGED
@@ -5,8 +5,8 @@ from utils.utils import *
5
 
6
 
7
  def detect(save_img=False):
8
- out, source, weights, half, view_img, save_txt, imgsz = \
9
- opt.output, opt.source, opt.weights, opt.half, opt.view_img, opt.save_txt, opt.img_size
10
  webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')
11
 
12
  # Initialize
@@ -14,7 +14,7 @@ def detect(save_img=False):
14
  if os.path.exists(out):
15
  shutil.rmtree(out) # delete output folder
16
  os.makedirs(out) # make new output folder
17
- half &= device.type != 'cpu' # half precision only supported on CUDA
18
 
19
  # Load model
20
  google_utils.attempt_download(weights)
@@ -142,7 +142,6 @@ if __name__ == '__main__':
142
  parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')
143
  parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
144
  parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
145
- parser.add_argument('--half', action='store_true', help='half precision FP16 inference')
146
  parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
147
  parser.add_argument('--view-img', action='store_true', help='display results')
148
  parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
 
5
 
6
 
7
  def detect(save_img=False):
8
+ out, source, weights, view_img, save_txt, imgsz = \
9
+ opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
10
  webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')
11
 
12
  # Initialize
 
14
  if os.path.exists(out):
15
  shutil.rmtree(out) # delete output folder
16
  os.makedirs(out) # make new output folder
17
+ half = device.type != 'cpu' # half precision only supported on CUDA
18
 
19
  # Load model
20
  google_utils.attempt_download(weights)
 
142
  parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')
143
  parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
144
  parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
 
145
  parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
146
  parser.add_argument('--view-img', action='store_true', help='display results')
147
  parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
test.py CHANGED
@@ -17,7 +17,6 @@ def test(data,
17
  save_json=False,
18
  single_cls=False,
19
  augment=False,
20
- half=False, # FP16
21
  model=None,
22
  dataloader=None,
23
  fast=False,
@@ -25,7 +24,7 @@ def test(data,
25
  # Initialize/load model and set device
26
  if model is None:
27
  device = torch_utils.select_device(opt.device, batch_size=batch_size)
28
- half &= device.type != 'cpu' # half precision only supported on CUDA
29
 
30
  # Remove previous
31
  for f in glob.glob('test_batch*.jpg'):
@@ -48,7 +47,8 @@ def test(data,
48
  device = next(model.parameters()).device # get model device
49
  training = True
50
 
51
- # Configure run
 
52
  with open(data) as f:
53
  data = yaml.load(f, Loader=yaml.FullLoader) # model dict
54
  nc = 1 if single_cls else int(data['nc']) # number of classes
@@ -57,7 +57,10 @@ def test(data,
57
  niou = iouv.numel()
58
 
59
  # Dataloader
60
- if dataloader is None:
 
 
 
61
  fast |= conf_thres > 0.001 # enable fast mode
62
  path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
63
  dataset = LoadImagesAndLabels(path,
@@ -75,9 +78,6 @@ def test(data,
75
  collate_fn=dataset.collate_fn)
76
 
77
  seen = 0
78
- model.eval()
79
- img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
80
- _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
81
  names = model.names if hasattr(model, 'names') else model.module.names
82
  coco91class = coco80_to_coco91_class()
83
  s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', '[email protected]', '[email protected]:.95')
@@ -221,11 +221,11 @@ def test(data,
221
  cocoDt = cocoGt.loadRes(f) # initialize COCO pred api
222
 
223
  cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
224
- cocoEval.params.imgIds = imgIds # [:32] # only evaluate these images
225
  cocoEval.evaluate()
226
  cocoEval.accumulate()
227
  cocoEval.summarize()
228
- map, map50 = cocoEval.stats[:2] # update to pycocotools results ([email protected]:0.95, [email protected])
229
  except:
230
  print('WARNING: pycocotools must be installed with numpy==1.17 to run correctly. '
231
  'See https://github.com/cocodataset/cocoapi/issues/356')
@@ -248,7 +248,6 @@ if __name__ == '__main__':
248
  parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
249
  parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
250
  parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
251
- parser.add_argument('--half', action='store_true', help='half precision FP16 inference')
252
  parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
253
  parser.add_argument('--augment', action='store_true', help='augmented inference')
254
  parser.add_argument('--verbose', action='store_true', help='report mAP by class')
@@ -268,8 +267,7 @@ if __name__ == '__main__':
268
  opt.iou_thres,
269
  opt.save_json,
270
  opt.single_cls,
271
- opt.augment,
272
- opt.half)
273
 
274
  elif opt.task == 'study': # run over a range of settings and save/plot
275
  for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
 
17
  save_json=False,
18
  single_cls=False,
19
  augment=False,
 
20
  model=None,
21
  dataloader=None,
22
  fast=False,
 
24
  # Initialize/load model and set device
25
  if model is None:
26
  device = torch_utils.select_device(opt.device, batch_size=batch_size)
27
+ half = device.type != 'cpu' # half precision only supported on CUDA
28
 
29
  # Remove previous
30
  for f in glob.glob('test_batch*.jpg'):
 
47
  device = next(model.parameters()).device # get model device
48
  training = True
49
 
50
+ # Configure
51
+ model.eval()
52
  with open(data) as f:
53
  data = yaml.load(f, Loader=yaml.FullLoader) # model dict
54
  nc = 1 if single_cls else int(data['nc']) # number of classes
 
57
  niou = iouv.numel()
58
 
59
  # Dataloader
60
+ if dataloader is None: # not training
61
+ img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
62
+ _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
63
+
64
  fast |= conf_thres > 0.001 # enable fast mode
65
  path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
66
  dataset = LoadImagesAndLabels(path,
 
78
  collate_fn=dataset.collate_fn)
79
 
80
  seen = 0
 
 
 
81
  names = model.names if hasattr(model, 'names') else model.module.names
82
  coco91class = coco80_to_coco91_class()
83
  s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', '[email protected]', '[email protected]:.95')
 
221
  cocoDt = cocoGt.loadRes(f) # initialize COCO pred api
222
 
223
  cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
224
+ cocoEval.params.imgIds = imgIds # image IDs to evaluate
225
  cocoEval.evaluate()
226
  cocoEval.accumulate()
227
  cocoEval.summarize()
228
+ map, map50 = cocoEval.stats[:2] # update results ([email protected]:0.95, [email protected])
229
  except:
230
  print('WARNING: pycocotools must be installed with numpy==1.17 to run correctly. '
231
  'See https://github.com/cocodataset/cocoapi/issues/356')
 
248
  parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
249
  parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
250
  parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
 
251
  parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
252
  parser.add_argument('--augment', action='store_true', help='augmented inference')
253
  parser.add_argument('--verbose', action='store_true', help='report mAP by class')
 
267
  opt.iou_thres,
268
  opt.save_json,
269
  opt.single_cls,
270
+ opt.augment)
 
271
 
272
  elif opt.task == 'study': # run over a range of settings and save/plot
273
  for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: