我是靠谱客的博主 饱满斑马,最近开发中收集的这篇文章主要介绍gather torch_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这种数值,而且为啥会在通道上进行划分?

经过一步步检查参数尺度,终于发现了问题:

4790c0acd91f6138057cf0d622fc881f.png

原来是这里的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源码:

a780638c599f44da88d08725a54f4a20.png

函数首先检测GPU数目,如果大于1的话,会用self.scatter将输入进行拆分:

032da7afeea9c00a7241ad8cff6bfea2.png

这里最终用到的函数是torch.cuda.comm.scatter:

b4a94df975ccf9e34269eb8d9f187568.png

所以和torch.scatter没什么区别的样机。

那处理起来就简单:

4433af8c71ec0dd64def1d1d3e56c3df.png

在定义输入变量的时候,先检测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的问题:维度不匹配所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错,欢迎将靠谱客网站推荐给程序员好友。

本图文内容来源于网友提供,作为学习参考使用,或来自网络收集整理,版权属于原作者所有。
点赞(51)

评论列表共有 0 条评论

立即
投稿
返回
顶部