RNN与torch DataParallel的爱恨情仇

BodyCsoulN

已于 2022-03-11 20:41:48 修改

阅读量1.6k

点赞数 1

文章标签： rnn lstm 深度学习 pytorch 人工智能

于 2022-03-11 20:40:07 首次发布

本文链接：https://blog.csdn.net/BodyCsoulN/article/details/123428447

版权

本文探讨了在使用PyTorch的DataParallel进行多GPU加速RNN模型时遇到的问题，主要涉及RNN的state向量在不同维度的处理，以及如何协调DataParallel的切分规则。作者指出，由于RNN层期望batch size在第二个维度，而DataParallel默认按第一维度切分，这可能导致错误。解决方案包括在state送入模型前将其batch size调整到第一个维度，计算前后进行维度转换，并在模型输出后恢复维度，以确保正确计算损失。

摘要由CSDN通过智能技术生成

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,</

最低0.47元/天解锁文章