RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0

记录一下遇到的改进网络模型时的错误。

起因:将CCNet的十字交叉注意力模块加入YOLOv5中,在写代码时,有部分方法用到了.cuda()这个方法来返回一个临时数据。YOLOv5在构建模型时,会跑一遍模型,此时数据类型是cpu类型的,当与cuda类型的临时数据进行计算时,自然会发生冲突。如果直接去掉.cuda()方法,临时数据变为cpu类型;虽然可以在第一次构建模型时不发生冲突,但是实际训练的时候,传入的数据是cuda类型的,因此会再次发生冲突。

解决方案:在传入的数据和临时数据计算之前,先将临时数据的类型转换为传入数据的类型就可以了。

def INF(B, H, W):
    return -torch.diag(torch.tensor(float("inf")).repeat(H), 0).unsqueeze(0).repeat(B * W, 1, 1)
    # return -torch.diag(torch.tensor(float("inf")).cuda().repeat(H), 0).unsqueeze(0).repeat(B * W, 1, 1)

......


energy_H = (torch.bmm(query_H, key_H) + self.INF(b2, h1, w1).to(query_H.device)).view(b2, w1, h1, h1).permute(0, 2, 1, 3)

如上面的代码所示,解决类型冲突的方法,就是令self.INF().to(传入数据.device)

你可能感兴趣的:(深度学习,python,人工智能)