PyTorch DataParallel的一些坑点和使用技巧

1. 报错PyTorch DataParallel RuntimeError: expected device cuda:1 but got device cuda:0

原因是直接使用了DataParallel作为神经网络。

# 可能报错的写法
model = DataParallel(Res_Net())
out = model(input)

# 解决办法
model = DataParallel(Res_Net())
model = model.modukle()
out = model(input)

2.关于模型的保存和使用

虽然我们可以直接保存DataParallel, 在load的时候直接对DataParallel进行load，但是这样的话模型就不能转移到单GPu上使用。所以我们统一对module部分load和save，这样模型既可以在单GPU也可以在多GPU上使用。

# save: 我们保存的是DataParallel的module（也就是我们真正要使用的神经网络）
model = resnet101()
model = torch.nn.DataParallel(model)
model = model.module()
torch.save(model.state_dict(), 'state')

# load： 我们对.module进行load， 直接对DataParallel进行load的话会报错，因为网络层的名称不同（parallel的网络层名称多了个.module）
model2 = resnet101()
model2 = torch.nn.DataParallel(model2)
model2 = model2.module()
model2.load_state_dict(torch.load('state'))