torch.nn.DataParallel
是通常用来多gpu加速的一个torch包。在我的实验过程中,发现了很多很多的bug。
在训练RNN模型时,由于RNN模型一般将batch_size
放在第二个维度(输入、state)。DataParallel
会将一个batch的数据切分开来,然后放到多个gpu上,默认按照第一维度进行切分(可以通过dim进行更换)。这里的问题就在于:
- 输入X是我们自己定义的,所以维度可以进行控制。
- state向量一般是由RNN layer自己生成,所以batch_size必然在第二个维度,(有一篇博客中指出,
batch_first=True
在一定程度上无效)。所以在DataParallel
进行切分的时候(数据从训练函数送入模型中),state向量会按照第一个维度(也是就序列长度)进行切分,然后在模型计算完成之后进行合并(从模型中返回到训练函数),这样,就与DataParallel
的初衷违背了,DataParallel
会将一个batch的数据切分开来。所以我们需要在这中间做个操作。
在tensor输入维度上可以选择第一位是batch size,或者第二位是batch size。理论上说这是个人习惯问题,只要前后统一就可以;但是在pytorch中内置的lstm层只接受batch second 的hidden层tensor。虽然在lstm层或padding层上可以定义batch_first=True,但是这只定义了输入tensor;hidden层仍然必须是batch second 格式。
————————————————
引用自CSDN博主「Edward Tivrusky IV」的原创文章,遵循CC 4.0 by-sa版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/yuuyuhaksho/article/details/87560640
————————————————
版权声明:本文为CSDN博主「pyxiea」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/xpy870663266/article/details/100002592
那我们应该怎么做呢?
- 在state从训练函数送入模型之前,将
batch_size
放在第一个维度,以使DataParallel
将state正常切分 - state在模型中开始计算之前,将
batch_size
放在第二个维度,以正常参与计算 - state参与计算之后,将
batch_size
放在第一个维度,以使多个gpu的state正常合并。 - state返回模型之后,将
batch_size
放在第二个维度,正常参与下一次计算。
def train_epoch():
...
state = state.permute([1,</