glenn-jocher commited on
Commit
be9edff
1 Parent(s): a18efc3

Update test.py --task speed and study (#2099)

Browse files

* Add --speed benchmark

* test range 256 - 1536

* update

* update

* update

* update

Files changed (2) hide show
  1. test.py +10 -6
  2. utils/plots.py +1 -1
test.py CHANGED
@@ -320,16 +320,20 @@ if __name__ == '__main__':
320
  save_conf=opt.save_conf,
321
  )
322
 
 
 
 
 
323
  elif opt.task == 'study': # run over a range of settings and save/plot
324
- for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
325
- f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem) # filename to save to
326
- x = list(range(320, 800, 64)) # x axis
327
  y = [] # y axis
328
  for i in x: # img-size
329
- print('\nRunning %s point %s...' % (f, i))
330
- r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,
331
  plots=False)
332
  y.append(r + t) # results and times
333
  np.savetxt(f, y, fmt='%10.4g') # save
334
  os.system('zip -r study.zip study_*.txt')
335
- plot_study_txt(f, x) # plot
 
320
  save_conf=opt.save_conf,
321
  )
322
 
323
+ elif opt.task == 'speed': # speed benchmarks
324
+ for w in opt.weights:
325
+ test(opt.data, w, opt.batch_size, opt.img_size, 0.25, 0.45, save_json=False, plots=False)
326
+
327
  elif opt.task == 'study': # run over a range of settings and save/plot
328
+ x = list(range(256, 1536 + 128, 128)) # x axis (image sizes)
329
+ for w in opt.weights:
330
+ f = f'study_{Path(opt.data).stem}_{Path(w).stem}.txt' # filename to save to
331
  y = [] # y axis
332
  for i in x: # img-size
333
+ print(f'\nRunning {f} point {i}...')
334
+ r, _, t = test(opt.data, w, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,
335
  plots=False)
336
  y.append(r + t) # results and times
337
  np.savetxt(f, y, fmt='%10.4g') # save
338
  os.system('zip -r study.zip study_*.txt')
339
+ plot_study_txt(x=x) # plot
utils/plots.py CHANGED
@@ -223,7 +223,7 @@ def plot_targets_txt(): # from utils.plots import *; plot_targets_txt()
223
  plt.savefig('targets.jpg', dpi=200)
224
 
225
 
226
- def plot_study_txt(path='study/', x=None): # from utils.plots import *; plot_study_txt()
227
  # Plot study.txt generated by test.py
228
  fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)
229
  ax = ax.ravel()
 
223
  plt.savefig('targets.jpg', dpi=200)
224
 
225
 
226
+ def plot_study_txt(path='', x=None): # from utils.plots import *; plot_study_txt()
227
  # Plot study.txt generated by test.py
228
  fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)
229
  ax = ax.ravel()