使用Pytorch多GPU训练RNN网络踩坑记录

问题介绍

在构建LAS端到端语音识别网络模型时，encoder使用了nn.GRU，decoder使用了nn.LSTMCell，在单个GPU上训练时，一切正常，使用nn.DataParallel进行多GPU训练时，一直报错。

问题

在使用DataParallel时，pytorch会把forward函数参数里同一个batch的数据拆分到不同GPU里，而不会拆分形如http://self.xxx的类的属性拆分掉。如果检查属性量的device，可以发现它们初始化在cuda：0中后没有改变设备号。

解决方案

1. 方案一：所有需要拆分的量，在forward函数中作为形参与返回值。在RNN中，要注意的问题还有一些。比如为了解决这个问题，隐藏状态常常被optimizer每次传入RNN的forward。此时，隐藏状态的初始化要在每一个epoch的开始，而不可以和optimizer在一起。这一细节问题网络上有的资料没有避免。
2. 方案二：不要同时使用，两者可以相互替代，只使用nn.LSTM、nn.GRU或者nn.LSTMCell、nn.GRUCell。