Commit
·
157aff2
1
Parent(s):
c6fc424
FP16 inference fix
Browse files
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
|
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']:
|