- 问题描述:
pytorch1.0,从pretrainedmodels 中import 了 resnet50模型
单GPU训练没有问题
但是改为多GPU训练的时候会报weight 和input不在一个device上的错误
- 解决方案
加入如下代码就可以了呢!!:
class Net(nn.Module):
def __init__(self, model):
super(Net, self).__init__()
self.l1 = nn.Sequential(*list(model.children())[:-1]).to('cuda:0')
self.last = list(model.children())[-1]
def forward(self, x):
x = self.l1(x)
x = x.view(x.size()[0], -1)
x = self.last(x)
return x
链接:https://github.com/Cadene/pretrained-models.pytorch/issues/120
或者可以从torchvision import
- 具体原因
也有大佬分析了哈哈哈哈哈
https://blog.csdn.net/senius/article/details/96599955