mirrors/yolov5
1
0
Fork 0
mirror of https://github.com/ultralytics/yolov5.git synced 2025-06-03 14:49:29 +08:00
Code Issues Projects Releases Wiki Activity

Add PyTorch-only benchmark arg (#7564)

Browse Source 
Usage:
```
python utils/benchmarks.py --weights yolov5s.pt --img 640 --pt-only
```
This commit is contained in:
Glenn Jocher 2022-04-24 20:03:23 -07:00 committed by GitHub
parent 950a85d9f6
commit e54e758cfb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
utils/benchmarks.py
View File

@ -53,6 +53,7 @@ def run(
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
half=False, # use FP16 half-precision inference half=False, # use FP16 half-precision inference
test=False, # test exports only test=False, # test exports only
pt_only=False, # test PyTorch only
): ):
y, t = [], time.time() y, t = [], time.time()
formats = export.export_formats() formats = export.export_formats()
@ -79,6 +80,8 @@ def run(
except Exception as e: except Exception as e:
LOGGER.warning(f'WARNING: Benchmark failure for {name}: {e}') LOGGER.warning(f'WARNING: Benchmark failure for {name}: {e}')
y.append([name, None, None]) # mAP, t_inference y.append([name, None, None]) # mAP, t_inference
if pt_only and i == 0:
break # break after PyTorch
# Print results # Print results
LOGGER.info('\n') LOGGER.info('\n')
@ -98,6 +101,7 @@ def test(
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
half=False, # use FP16 half-precision inference half=False, # use FP16 half-precision inference
test=False, # test exports only test=False, # test exports only
pt_only=False, # test PyTorch only
): ):
y, t = [], time.time() y, t = [], time.time()
formats = export.export_formats() formats = export.export_formats()
@ -130,6 +134,7 @@ def parse_opt():
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference') parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
parser.add_argument('--test', action='store_true', help='test exports only') parser.add_argument('--test', action='store_true', help='test exports only')
parser.add_argument('--pt-only', action='store_true', help='test PyTorch only')
opt = parser.parse_args() opt = parser.parse_args()
print_args(vars(opt)) print_args(vars(opt))
return opt return opt