glenn-jocher commited on
Commit
e54e758
·
unverified ·
1 Parent(s): 950a85d

Add PyTorch-only benchmark arg (#7564)

Browse files

Usage:
```
python utils/benchmarks.py --weights yolov5s.pt --img 640 --pt-only
```

Files changed (1) hide show
  1. 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