YOLO v5/v8实例分割报错

最近在使用yolov8跑实例分割的时候，发现了以下错误。

Traceback (most recent call last):
  File "e:\VS Code Document\PyTorch\YOLO\ultralytics-main\main_segment.py", line 117, in 
    predict()
  File "e:\VS Code Document\PyTorch\YOLO\ultralytics-main\main_segment.py", line 113, in predict
    predictor.predict_cli()
  File "e:\VS Code Document\PyTorch\YOLO\ultralytics-main\ultralytics\yolo\engine\predictor.py", line 120, in predict_cli
    for _ in gen:  # running CLI inference without accumulating any outputs (do not modify)
  File "F:\python3.10\lib\site-packages\torch\autograd\grad_mode.py", line 43, in generator_context
    response = gen.send(None)
  File "e:\VS Code Document\PyTorch\YOLO\ultralytics-main\ultralytics\yolo\engine\predictor.py", line 187, in stream_inference
    s += self.write_results(i, self.results, (p, im, im0))
  File "e:\VS Code Document\PyTorch\YOLO\ultralytics-main\main_segment.py", line 74, in write_results
    self.annotator.masks(masks=mask.masks, colors=[colors(x, True) for x in det.cls], im_gpu=im_gpu)
  File "e:\VS Code Document\PyTorch\YOLO\ultralytics-main\ultralytics\yolo\utils\plotting.py", line 121, in masks
    masks_color = masks * (colors * alpha)  # shape(n,h,w,3)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

读一下大概意思就是：期望tensor运算在相同的device上，但是至少发现了2个device。 

直接定位报错地方在plotting.py的 masks_color = masks * (colors * alpha)  # shape(n,h,w,3)处二话不说，先把参与运算的3个变量类型打印出来。

print(type(masks))
print(type(colors))
print(type(alpha))
exit()

发现输出

 继续把两个tensor的device位置打印出来

print(masks.device)
print(colors.device)
exit()

发现masks在cpu上，而colors在cuda:0上，这样问题就很明了了。

解决方法

将

masks = masks.unsqueeze(3)

改为即可

masks = masks.unsqueeze(3).to(im_gpu.device) # shape(n,h,w,1)

其实认真读一下前面的代码，就会发现，mask在传入和unsqueeze操作的时候，没有放在GPU上。

PS：上述方法只对Ultralytics YOLOv8.0.44有效，其他版本不保证有效。