Hugging Face载入模型双卡训练报错

Question

模型载入代码如下:

device = torch.device("cuda:0")
model6 = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny",
                                                       ignore_mismatched_sizes=True)  # /upernet-swin-large,upernet-convnext-tiny
model6 = nn.DataParallel(model6, device_ids = [0, 1])
model6 = model6.to(device)
x2 = torch.randn(4, 3, 256,256).to(device)
print(model6(x2).shape)

报错:

RuntimeError: Expected tensor for argument #1 'input' to have the same device as
tensor for argument #2 'weight'; but device 1 does not equal 0 (while checking 
arguments for cudnn_convolution)

请问如何解决?

Answer 1

模型的输入数据和模型权重必须在同一设备上。

代码里模型移动到了GPU（device 0）上，执行 model6(x2) 时，输入数据 x2 却没有明确地分配到某一个GPU上。

试试使用 nn.DataParallel 的 forward() 方法来解决这个问题。将 model6(x2) 更改为 model6.module.forward(x2)，看看能否解决问题。这样会确保你的输入数据被分配到所有的GPU上。

device = torch.device("cuda:0")
model6 = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny",
                                                       ignore_mismatched_sizes=True)  # /upernet-swin-large,upernet-convnext-tiny
model6 = nn.DataParallel(model6, device_ids=[0, 1])
model6 = model6.to(device)
x2 = torch.randn(4, 3, 256, 256).to(device)
print(model6.module.forward(x2).shape)