Add PyTorch-only benchmark arg (#7564)
Browse filesUsage:
```
python utils/benchmarks.py --weights yolov5s.pt --img 640 --pt-only
```
- utils/benchmarks.py +5 -0
utils/benchmarks.py
CHANGED
@@ -53,6 +53,7 @@ def run(
|
|
53 |
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
54 |
half=False, # use FP16 half-precision inference
|
55 |
test=False, # test exports only
|
|
|
56 |
):
|
57 |
y, t = [], time.time()
|
58 |
formats = export.export_formats()
|
@@ -79,6 +80,8 @@ def run(
|
|
79 |
except Exception as e:
|
80 |
LOGGER.warning(f'WARNING: Benchmark failure for {name}: {e}')
|
81 |
y.append([name, None, None]) # mAP, t_inference
|
|
|
|
|
82 |
|
83 |
# Print results
|
84 |
LOGGER.info('\n')
|
@@ -98,6 +101,7 @@ def test(
|
|
98 |
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
99 |
half=False, # use FP16 half-precision inference
|
100 |
test=False, # test exports only
|
|
|
101 |
):
|
102 |
y, t = [], time.time()
|
103 |
formats = export.export_formats()
|
@@ -130,6 +134,7 @@ def parse_opt():
|
|
130 |
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
131 |
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
132 |
parser.add_argument('--test', action='store_true', help='test exports only')
|
|
|
133 |
opt = parser.parse_args()
|
134 |
print_args(vars(opt))
|
135 |
return opt
|
|
|
53 |
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
54 |
half=False, # use FP16 half-precision inference
|
55 |
test=False, # test exports only
|
56 |
+
pt_only=False, # test PyTorch only
|
57 |
):
|
58 |
y, t = [], time.time()
|
59 |
formats = export.export_formats()
|
|
|
80 |
except Exception as e:
|
81 |
LOGGER.warning(f'WARNING: Benchmark failure for {name}: {e}')
|
82 |
y.append([name, None, None]) # mAP, t_inference
|
83 |
+
if pt_only and i == 0:
|
84 |
+
break # break after PyTorch
|
85 |
|
86 |
# Print results
|
87 |
LOGGER.info('\n')
|
|
|
101 |
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
102 |
half=False, # use FP16 half-precision inference
|
103 |
test=False, # test exports only
|
104 |
+
pt_only=False, # test PyTorch only
|
105 |
):
|
106 |
y, t = [], time.time()
|
107 |
formats = export.export_formats()
|
|
|
134 |
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
135 |
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
136 |
parser.add_argument('--test', action='store_true', help='test exports only')
|
137 |
+
parser.add_argument('--pt-only', action='store_true', help='test PyTorch only')
|
138 |
opt = parser.parse_args()
|
139 |
print_args(vars(opt))
|
140 |
return opt
|