使用的是yolov8 pytorch1.12.0+cuda11.3,出现了这种情况,已经设置了device使用GPU,训练时还是会提示不能使用CPU报错,根据网上MissingName博主的回答,解决办法是:
1.使用GPU运行,之后似乎会生成一个全0张量将整个模型跑一次,而这个全0张量默认是在CPU中运行的,因为不知道哪一个会起效而且计算时要在同一个设备运行,于是我将模型\ultralytics-main\ultralytics\nn\task.py下的class DetectionModel(BaseModel):中的增加一些内容:
if isinstance(m, (Detect, Segment, Pose)):
s = 256 # 2x min stride
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")#增加的,将模型放入GPU运行
m.inplace = self.inplace
m.cuda() #增加的,将模型放入GPU运行
forward = lambda x: self.forward(x)[0] if isinstance(m, (Segment, Pose)) else self.forward(x)
m.stride = torch.tensor([s / x.shape[-2] for x in forward(torch.zeros(1, ch, s, s).to(device))]) # forward 在张量后加了一个to(device)
self.stride = m.stride
m.bias_init() # only run once
2.生成的全0张量和model都放入GPU运算。第一次遍历模型使用的def _predict_once方法中也加入self.model.cuda()
def _predict_once(self, x, profile=False, visualize=False, embed=None):
self.model.cuda()
之后就可以成功运行了!