mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

NMS MPS device wrapper (#9620) 

* NMS MPS device wrapper

May resolve https://github.com/ultralytics/yolov5/issues/9613

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

* Update general.py

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/9558/head^2
Glenn Jocher 2022-09-27 18:02:48 +02:00 committed by GitHub
parent 7314363f26
commit 2373d5470e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
utils/general.py
View File

@ -843,7 +843,9 @@ def non_max_suppression(
if isinstance(prediction, (list, tuple)): # YOLOv5 model in validation model, output = (inference_out, loss_out)
prediction = prediction[0] # select only inference output
if 'mps' in prediction.device.type: # MPS not fully supported yet, convert tensors to CPU before NMS
device = prediction.device
mps = 'mps' in device.type # Apple MPS
if mps: # MPS not fully supported yet, convert tensors to CPU before NMS
prediction = prediction.cpu()
bs = prediction.shape[0] # batch size
nc = prediction.shape[2] - nm - 5 # number of classes
@ -930,6 +932,8 @@ def non_max_suppression(
i = i[iou.sum(1) > 1] # require redundancy
output[xi] = x[i]
if mps:
output[xi] = output[xi].to(device)
if (time.time() - t) > time_limit:
LOGGER.warning(f'WARNING ⚠️ NMS time limit {time_limit:.3f}s exceeded')
break # time limit exceeded