如何在 Pytorch 中使用 DataParallel 并行化 RNN 函数
Posted
技术标签:
【中文标题】如何在 Pytorch 中使用 DataParallel 并行化 RNN 函数【英文标题】:How to parallelize RNN function in Pytorch with DataParallel 【发布时间】:2017-11-19 14:10:32 【问题描述】:这是一个运行基于字符的语言生成的 RNN 模型:
class RNN(nn.Module):
def __init__(self, input_size, hidden_size, output_size, n_layers):
super(RNN, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.output_size = output_size
self.n_layers = n_layers
self.encoder = nn.Embedding(input_size, hidden_size)
self.GRU = nn.GRU(hidden_size, hidden_size, n_layers, batch_first=True)
self.decoder = nn.Linear(hidden_size, output_size)
def forward(self, input, batch_size):
self.init_hidden(batch_size)
input = self.encoder(input)
output, self.hidden = self.GRU(input, self.hidden)
output = self.decoder(output.view(batch_size, self.hidden_size))
return output
def init_hidden(self, batch_size):
self.hidden = Variable(torch.randn(self.n_layers, batch_size, self.hidden_size).cuda())
我使用 DataParallel 实例化模型,以便在我的 4 个 GPU 上拆分输入批次:
net = torch.nn.DataParallel(RNN(n_chars, hidden_size, n_chars, n_layers)).cuda()
这是full code。
不幸的是,DataParallel 要求输入具有 batch_size 作为第一维,但 GRU 函数要求隐藏张量具有 batch_size 作为第二维:
output, self.hidden = self.GRU(input, self.hidden)
原样的代码会引发以下错误(请注意显示编码器在 4 个 GPU 上正确执行的打印输出):
...
forward function: encoding input of shape: (16L, 1L)
forward function: encoding input of shape: (16L, 1L)
forward function: encoding input of shape: (16L,
forward function: encoding input of shape:
forward function: GRU processing input of shape:
1L)
( (16L, 16L1L, 1L), 100L)
forward function: GRU processing input of shape:
(16L, 1L,
forward function: GRU processing input of shape:100L)
(16L
forward function: GRU processing input of shape:, 1L, 100L) (
16L, 1L, 100L)
Traceback (most recent call last):
File "gru2.py", line 166, in <module>
output = net(c, batch_size)
File "/root/miniconda2/lib/python2.7/site-packages/torch/nn/modules/module.py", line 206, in __call__
result = self.forward(*input, **kwargs)
File "/root/miniconda2/lib/python2.7/site-packages/torch/nn/parallel/data_parallel.py", line 61, in forward
outputs = self.parallel_apply(replicas, inputs, kwargs)
File "/root/miniconda2/lib/python2.7/site-packages/torch/nn/parallel/data_parallel.py", line 71, in parallel_apply
return parallel_apply(replicas, inputs, kwargs)
File "/root/miniconda2/lib/python2.7/site-packages/torch/nn/parallel/parallel_apply.py", line 45, in parallel_apply
raise output
RuntimeError: Expected hidden size (2, 16L, 100), got (2L, 64L, 100L)
这里模型有2层,batch_size=64,hidden_size=100。
如何在 forward 函数中并行化 GRU 操作?
【问题讨论】:
【参考方案1】:您可以简单地设置参数dim=1,例如
net = torch.nn.DataParallel(RNN(n_chars, hidden_size, n_chars, n_layers), dim=1).cuda()
【讨论】:
【参考方案2】:PyTorch 1.5 完全解决了 RNN 训练和 DataParallel 的问题。 看起来它做得非常无缝。 不再需要 gerrymandering。 我今天在一个涉及双向 GRUS 语音 mfcc 的项目中证实了这一点。
class PEncoder(nn.Module):
def __init__(self, args, encoder):
super(PEncoder, self).__init__()
self.gpu_ids = args.gpu_ids
self.model = encoder
def forward(self, input):
if len(self.gpu_ids) > 1:
return nn.parallel.data_parallel(self.model, (input), self.gpu_ids)
else:
return self.model(input)
就是这么简单。 这确实将您的模型包装在另一个模型中,并有效地产生了稍微不同的计算图。因此,如果您有较早训练过的模型,您可能必须以特殊方式加载它们并为此并行包装创建一些设置器。试一试,你会看到的。 (我还没有确认这方面的内容。)
【讨论】:
这也可以帮助pytorch.org/docs/stable/notes/…以上是关于如何在 Pytorch 中使用 DataParallel 并行化 RNN 函数的主要内容,如果未能解决你的问题,请参考以下文章
如何在 pytorch 和 tensorflow 中使用张量核心?
如何在 PyTorch 中使用 autograd.gradcheck?