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

FP16 inference fix

Browse files
Files changed (1) hide show
  1. test.py +5 -4
test.py CHANGED
@@ -13,15 +13,15 @@ def test(data,
13
  batch_size=16,
14
  imgsz=640,
15
  conf_thres=0.001,
16
- iou_thres=0.6, # for nms
17
  save_json=False,
18
  single_cls=False,
19
  augment=False,
 
20
  model=None,
21
  dataloader=None,
22
  fast=False,
23
- verbose=False,
24
- half=False): # FP16
25
  # Initialize/load model and set device
26
  if model is None:
27
  device = torch_utils.select_device(opt.device, batch_size=batch_size)
@@ -268,7 +268,8 @@ if __name__ == '__main__':
268
  opt.iou_thres,
269
  opt.save_json,
270
  opt.single_cls,
271
- opt.augment)
 
272
 
273
  elif opt.task == 'study': # run over a range of settings and save/plot
274
  for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
 
13
  batch_size=16,
14
  imgsz=640,
15
  conf_thres=0.001,
16
+ iou_thres=0.6, # for NMS
17
  save_json=False,
18
  single_cls=False,
19
  augment=False,
20
+ half=False, # FP16
21
  model=None,
22
  dataloader=None,
23
  fast=False,
24
+ verbose=False):
 
25
  # Initialize/load model and set device
26
  if model is None:
27
  device = torch_utils.select_device(opt.device, batch_size=batch_size)
 
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']: