概述
Pytorch中的多GPU非常好用,一句话就能搞定:self.model = torch.nn.DataParallel(self.model)。
然而这两天我做零样本学习中遇到了一个问题始终无法解决,就是说单GPU可以跑,一旦使用多GPU,就会出现:
RuntimeError: Gather got an input of invalid size: got [24, 10, 448,448], but expected [24, 11, 448,448] (gather at /pytorch/torch/csrc/cuda/comm.cpp:239
的错误。我每个batch数据为48张448*448图片,在每个gpu上划分了24张没问题,但通道上不知道为啥会出现10/11这两个数字,我backbone用的是ResNet-101,哪一层都没有10或11这种数值,而且为啥会在通道上进行划分?
经过一步步检查参数尺度,终于发现了问题:
原来是这里的self.embed_arr导致的问题。self.embed_arr是我输入的语义标签,和批次无关,这个变量在整个训练过程中都是固定的,维度为21*300,会影响到特征通道。在单GPU情况下,对每张图片self.embed_arr都是一样的,因此正常。然而由于torch.nn.DataParallel(self.model),所以model的input都会被强制等分,因此当我使用两个gpu时,在gpu0上和gpu1上的self.embed_arr被分别拆成了10*300和11*300两份,于是我的图像特征也在两个gpu上变成了[24, 10, 448,448]和[24, 11, 448,448]。当两个gpu forward完就会发现数据维度不匹配,无法gather。终于破案了!(如果self.embed_arr的维度恰好为[20*300],在两个GPU划分为两个10*30,模型正常运行,细思极恐~)
接下来就是如何解决这个问题,我查看了torch.nn.DataParallel源码:
函数首先检测GPU数目,如果大于1的话,会用self.scatter将输入进行拆分:
这里最终用到的函数是torch.cuda.comm.scatter:
所以和torch.scatter没什么区别的样机。
那处理起来就简单:
在定义输入变量的时候,先检测gpu数目,如果是多GPU的话,就把输入变量复制几分,这样model.forward的时候,拆分出来在每个GPU上的就一样了。。。
除此之外,我还想到在定义model的时候,以成员变量的形式把embed_att保存在模型对象里,然后在forward函数里,直接调用self.成员,但实际发现会导致在GPU#1上跑的时候,参数和变量不在同一个GPU的问题。
这里想@pytorch,能不能在函数的输入那块加一个可以控制输入是否需要拆分的功能,哈哈哈。。。
最后
以上就是饱满斑马为你收集整理的gather torch_torch.nn.DataParallel中数据Gather的问题:维度不匹配的全部内容,希望文章能够帮你解决gather torch_torch.nn.DataParallel中数据Gather的问题:维度不匹配所遇到的程序开发问题。
如果觉得靠谱客网站的内容还不错,欢迎将靠谱客网站推荐给程序员好友。
发表评论 取消回复