解决了PyTorch 使用torch.nn.DataParallel 进行多GPU训练的一个BUG:模型(参数)和数据不在相同设备上

解决了PyTorch 使用torch.nn.DataParallel 进行多GPU训练的一个BUG:模型(参数)和数据不在相同设备上

博客地址：https://senyang-ml.github.io/2019/07/20/pytorch-multigpu/

解决了PyTorch 使用torch.nn.DataParallel 进行多GPU训练的一个BUG:

模型(参数)和数据不在相同设备上

使用torch.nn.DataParallel进行多GPU训练时出现了一个BUG, 困扰许久:

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)

这个错误表明, input数据在device 1 上, 而模型的参数在device 0 上 (暗示数据是被拆分到了各个GPU上,但是BUG出现位置的此处参数可能没有成功复制到其他GPU上, 或者说, 还是调用了复制前的那个参数地址)

之前调试了好久, 也没有解决掉, 在Github上有一个issue和我的问题很像:
https://github.com/pytorch/pytorch/issues/8637
但是我还是没有找到自己的问题在哪里.

今天, 我又准备再此挑战这个难题

经过6个小时的print调试法以及后面关键的VScode的Debug功能, 我大功告成,找到了问题所在！！！

在我的A(nn.Module)类的forward 前向计算函数里面, 有一处调用了一个该类的列表self.cell_fabrics属性, 其中，列表中的元素是直接通过

self.cell_fabrics = [self.cell_1,...,self.cell_n]

这个语句来赋值的, 其中每个self.cell也是nn.Module类。

当我们使用torch.nn.DataParallel复制了A(nn.Module)类的模型放到设备cuda:1上后（注意：cuda:0上始终保存着原始的模型），并且当多GPU并行处理的程序，运行到下面代码段时:

def forward(self,x)
    for layer in self.cell_fabrics:
        for cell in layer:
           y = cell (x,)

debug输出显示：x是在设备cuda:1 上面, 而 cell 中的参数明显都在 cuda:0上。

也就是说:

此时，前传调用的self.cell_fabrics 列表中的各个元素 （self.cell）的地址，还是指向在没有进行torch.nn.DataParallel之前的nn.module 的那些self.cell, 而nn.DataParallel类的nn.module的参数都默认存放在device(type='cuda',index=0)上 .

torch.nn.Module的广播机制和列表赋值方式

torch.nn.DataParallel(model,device_ids=[range(len(gpus))])的机制是, 将属于nn.module类的model以及其广播的所有nn.module子类的上的所有参数,复制成len(gpus)份,送到各个GPU上. 这种传播必须通过其子类通过注册(register)成为nn.module的属性才得以完成, 其属性为列表中的元素是不会被复制的, 所以其属性中的元素还是存放在默认设备device 0上

所以 在使用torch.nn.DataParallel进行多GPU训练的时候, 请注意：所有属于模型参数的模块以及其子模块必须以nn.Module的类型注册为模型的属性, 如果需要一个列表来批量存放子模块或者参数时, 请采用nn.ModuleList或者nn.ModuleDict这样的继承了nn.Module的类来进行定义, 并且在forward(self,)前向传播的过程中，需要直接调用属于 nn.Module,nn.ModuleList或者nn.ModuleDict 这样的属性。

那么torch.nn.DataParallel将会正常地将模型参数准确复制到多个GPU上, 并默认根据数据的batchsize的大小平分成GPU的数量分别送到相应的GPU设备上,

运用多线程的方式, 同时对这些数据进行加工处理, 然后收集各个GPU上最终产生对模型的参数的梯度, 汇总到一起更新原模型的参数!

参考:

1. https://github.com/pytorch/pytorch/issues/8637
2. https://pytorch.org/docs/stable/nn.html#dataparallel-layers-multi-gpu-distributed