Browse Source

FP16 inference fix

5.0
Glenn Jocher 4 years ago
parent
commit
157aff2854
1 changed files with 5 additions and 4 deletions
  1. +5
    -4
      test.py

+ 5
- 4
test.py View File

batch_size=16, batch_size=16,
imgsz=640, imgsz=640,
conf_thres=0.001, conf_thres=0.001,
iou_thres=0.6, # for nms
iou_thres=0.6, # for NMS
save_json=False, save_json=False,
single_cls=False, single_cls=False,
augment=False, augment=False,
half=False, # FP16
model=None, model=None,
dataloader=None, dataloader=None,
fast=False, fast=False,
verbose=False,
half=False): # FP16
verbose=False):
# Initialize/load model and set device # Initialize/load model and set device
if model is None: if model is None:
device = torch_utils.select_device(opt.device, batch_size=batch_size) device = torch_utils.select_device(opt.device, batch_size=batch_size)
opt.iou_thres, opt.iou_thres,
opt.save_json, opt.save_json,
opt.single_cls, opt.single_cls,
opt.augment)
opt.augment,
opt.half)


elif opt.task == 'study': # run over a range of settings and save/plot elif opt.task == 'study': # run over a range of settings and save/plot
for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']: for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:

Loading…
Cancel
Save