pytorch DistributedDataParallel单机多卡 pytorch多卡训练报错

自己的代码使用多GPU训练有一段时间了，直到最近，发现了一点点异常，虽然代码已经按照指定的GPU训练，但是，最近经别人提醒我才发现，其实代码在所有GPU上都占用了资源，不多，大概只有2M，但是有时候会产生一些意想不到的错误。所以准备记录一下所有的踩到的坑，并奉上一份demo，希望炼丹者们少浪费点时间在踩坑上。

千言万语汇成一段话：

os.environ["CUDA_VISIBLE_DEVICES"] = GPU_IDS

这行代码必不可少，之前的代码虽然能够达到多GPU训练的效果，但是因为没有这行代码，导致代码在所有的GPU都占用了大概2M的资源，这对于强迫症玩家来说很不友好（我指定你在第哪些张卡上跑 你就必须在第哪些张卡上跑！），如下图，这里我虽然指定了0，1，2卡，但是可以看到其它卡都有2M的占用，如果使用 fuser -v /dev/nvidia*命令，就可以看到，八张卡都占用了资源。

为了解决这个问题，后来又在代码里加上了上面这行代码，该行代码的作用是指定代码可以看到哪些GPU，没有指定的卡，代码可以认为不存在（差不多就是这么个意思），但是加载模型的时候，不能使用肉眼看到的这些卡的编号，而是使用映射编号（做了无数次实验才发现，查了很多博客没碰到这个问题）

这里的GPU_IDS为磁盘上的真实卡的编号，如果我想指定5，6，7最后三张卡训练我的模型，那么代码就应该是：os.environ["CUDA_VISIBLE_DEVICES"] = "5,6,7" 但是

torch.nn.DataParallel(model, device_ids=devices)

这个devices不能写[5,6,7]，而应该是[0,1,2]。代码只能看到三张卡，这三张卡的顺序从0开始编号。

实验完，总结一下，就这么简单。

下面记录 两个单机单卡和单机多卡训练的demo （多机多卡？ 穷啊…）

单机单卡

# 按照上面的逻辑，os.environ写在所有代码的开头，就假设编号为7的卡吧
os.environ["CUDA_VISIBLE_DEVICES"] = '7'

# 接着是device。这里的 cuda:0 不能变！必须是0！
# 注意！！！第一行代码和第二行代码一定按照顺序，如果颠倒顺序，就等着找bug吧
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# 加载模型
model = Unet_3D(1, 1).to(device)

# 数据加载
data = data.to(device)

# 如果你的损失函数是自定义的，那么也需要加载一下
criterion = DiceLoss().to(device)

# 保存模型
torch.save(model.state_dict(), save_path)

单机多卡

# 这里指定使用后面的三张卡训练,为了方便，我把代码写的健壮一点。
devices = [5,6,7] # 这个是要训练的卡，nvidia-smi显示出的肉眼可见的卡

# 设置os.environ
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(i) for i in devices])

# devices
device_ids = [i for i in range(len(devices))]

# 加载模型
model = Unet_3D(1, 1)
model = torch.nn.DataParallel(model, device_ids=device_ids)
model = model.cuda(device_ids[0])

# 数据加载
data = data.cuda(device=device_ids[0])

# 自定义的损失函数
criterion = DiceLoss().cuda(device=device_ids[0])

# 保存模型, 多GPU训练时，保存模型要加个module，否则测试的时候读取模型又是一言难尽
torch.save(model.module.state_dict(), save_path)

PS： 如果你多GPU训练模型时忘记使用 model.module.state_dict()而使用了model.state_dict(), 也不慌，网上有大佬给出了办法（链接我找不到了，如果原博看到还望给个链接）主要思想就是剥去外壳把模型的参数都读出来，赋值到新字典里。

def load_network(model_path):

    # Model was saved by mutiple GPUs training but tested by single GPU
    
    model = Unet_3D(1, 1).to(device)
    state_dict = torch.load(model_path, map_location='cuda:0')

    from collections import OrderedDict
    new_state_dict = OrderedDict()

    for k, v in state_dict.items():
        namekey = k[7:] # remove `module.`
        new_state_dict[namekey] = v
    
    model.load_state_dict(new_state_dict)
    
    return model

此外，有时候你使用了多GPU训练，但是发现速度好像并没有提升？？？，这也是遇到过的一个问题，产生这个问题的原因主要是在加载数据时一个参数的设置，

train_load = DataLoader(dataset=train_set, batch_size=args.batch_size*len(args.device_ids), shuffle=True, num_workers=8)

num_workers的设置对训练效率尤其重要，这个参数的设置需要根据自己的硬件设置选择。

---

能想到就这么多了，有问题欢迎交流。