您的位置:首页 > 其它

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

2019-07-20 16:34 8639 查看

解决了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_2,...,self.cell_n]
来赋值的,其中每个
self.cell
也是
nn.Module

即用

self.cell_fabrics = [self.cell_0_0, self.cell_0_1, … , self.cell_3_0, self.cell_3_1,…, self.cell_5_0]
这样的方式,将所有的
cell类
放到
A类
的一个
列表属性
中, 而当整个
A类
通过
torch.nn.DataParallel
被复制了一份放到设备
cuda:1
上时, 当该设备运行到下面代码段时:

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

x
是在设备
cuda:1
上面, 而 debug显示
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.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
内容来自用户分享和网络整理,不保证内容的准确性,如有侵权内容,可联系管理员处理 点击这里给我发消息
标签: 
相关文章推荐