FP16 inference fix
This commit is contained in:
parent
c6fc4242a2
commit
157aff2854
9
test.py
9
test.py
|
|
@ -13,15 +13,15 @@ def test(data,
|
||||||
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,
|
verbose=False):
|
||||||
half=False): # FP16
|
|
||||||
# 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)
|
||||||
|
|
@ -268,7 +268,8 @@ if __name__ == '__main__':
|
||||||
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…
Reference in New Issue