pytorch gather_torch.nn.DataParallel中数据Gather的问题：维度不匹配

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，能不能在函数的输入那块加一个可以控制输入是否需要拆分的功能，哈哈哈。。。