RNN与torch DataParallel的爱恨情仇

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

那我们应该怎么做呢？

1. 在state从训练函数送入模型之前，将batch_size放在第一个维度，以使DataParallel将state正常切分
2. state在模型中开始计算之前，将batch_size放在第二个维度，以正常参与计算
3. state参与计算之后，将batch_size放在第一个维度，以使多个gpu的state正常合并。
4. state返回模型之后，将batch_size放在第二个维度，正常参与下一次计算。

def train_epoch():
	...
	state = state.permute([1,</