greg2451 pre-commit-ci[bot] commited on
Commit
8fcdf3b
1 Parent(s): 079b36d

Fixing minor multi-streaming issues with TensoRT engine (#6504)

Browse files

* Update batch-size in model.warmup() + indentation for logging inference results

* These changes are in response to PR comments

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Files changed (2) hide show
  1. detect.py +4 -4
  2. val.py +1 -1
detect.py CHANGED
@@ -110,7 +110,7 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
110
  vid_path, vid_writer = [None] * bs, [None] * bs
111
 
112
  # Run inference
113
- model.warmup(imgsz=(1, 3, *imgsz), half=half) # warmup
114
  dt, seen = [0.0, 0.0, 0.0], 0
115
  for path, im, im0s, vid_cap, s in dataset:
116
  t1 = time_sync()
@@ -175,9 +175,6 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
175
  if save_crop:
176
  save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
177
 
178
- # Print time (inference-only)
179
- LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
180
-
181
  # Stream results
182
  im0 = annotator.result()
183
  if view_img:
@@ -203,6 +200,9 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
203
  vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
204
  vid_writer[i].write(im0)
205
 
 
 
 
206
  # Print results
207
  t = tuple(x / seen * 1E3 for x in dt) # speeds per image
208
  LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
 
110
  vid_path, vid_writer = [None] * bs, [None] * bs
111
 
112
  # Run inference
113
+ model.warmup(imgsz=(1 if pt else bs, 3, *imgsz), half=half) # warmup
114
  dt, seen = [0.0, 0.0, 0.0], 0
115
  for path, im, im0s, vid_cap, s in dataset:
116
  t1 = time_sync()
 
175
  if save_crop:
176
  save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
177
 
 
 
 
178
  # Stream results
179
  im0 = annotator.result()
180
  if view_img:
 
200
  vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
201
  vid_writer[i].write(im0)
202
 
203
+ # Print time (inference-only)
204
+ LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
205
+
206
  # Print results
207
  t = tuple(x / seen * 1E3 for x in dt) # speeds per image
208
  LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
val.py CHANGED
@@ -162,7 +162,7 @@ def run(data,
162
 
163
  # Dataloader
164
  if not training:
165
- model.warmup(imgsz=(1, 3, imgsz, imgsz), half=half) # warmup
166
  pad = 0.0 if task == 'speed' else 0.5
167
  task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images
168
  dataloader = create_dataloader(data[task], imgsz, batch_size, stride, single_cls, pad=pad, rect=pt,
 
162
 
163
  # Dataloader
164
  if not training:
165
+ model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz), half=half) # warmup
166
  pad = 0.0 if task == 'speed' else 0.5
167
  task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images
168
  dataloader = create_dataloader(data[task], imgsz, batch_size, stride, single_cls, pad=pad, rect=pt,