完整报错:
module must have its parameters and buffers on device cuda:2 (device_ids[0]) but found one of them on device: cuda:0
报错原因:
使用多 GPU 需要有一个主 GPU,来把每个 batch 的数据分发到每个 GPU,并从每个 GPU 收集计算好的结果。如果不指定主 GPU,那么数据就直接分发到每个 GPU,会造成有些数据在某个 GPU,而另一部分数据在其他 GPU,计算出错。
解决办法:
使用 DataParallel 时,device 指定某个 GPU 为 主 GPU:
net = nn.DataParallel(net)
device = torch.device("cuda:2" ) # 指定gpu2为主GPU
net.to(device)