在具有不同隐藏大小的PyTorch LSTM模型上附加一个递归层

时间:2018-12-08 19:18:16

标签: python lstm pytorch

我正在开发一个BI-LSTM模型,用于使用PyTorch进行序列分析。我正在为此使用torch.nn.LSTM。使用该模块,您可以通过仅传递参数num_layers作为层数(例如num_layers=2)而具有多个层。但是,所有这些对象都具有相同的hidden_size,这对我来说部分很好,我只想让所有的对象都具有相同的hidden_size 但是最后一层的大小不同。基本示例如下:

rnn = nn.LSTM(input_size=10, hidden_size=20, num_layers=2)
inp = torch.randn(5, 3, 10)
h0 = torch.randn(2, 3, 20)
c0 = torch.randn(2, 3, 20)
output, (hn, cn) = rnn(inp, (h0, c0))

输出暗淡为(5, 3, 20

一个解决方案(但对我不利)正在实现额外的模型,该模型输出我需要的尺寸并从第一个模型中获取输入,例如:

rnn_two = nn.LSTM(input_size=20, hidden_size=2)
output2, _ = rnn_two(output)

但是,我不想这样做,因为我使用DataParallel 并行化模型,所以我需要全部成为一个包。我希望找到类似于keras的东西,例如:

rnn.add(LSTM, hidden_size=2)

我检查了LSTM source code,但找不到我需要的东西。

有什么建议吗?

2 个答案:

答案 0 :(得分:2)

如果我没记错的话,可以这样做:

import torch.nn as nn
import torch.nn.functional as F

class RnnWith2HiddenSizesModel(nn.Module):
    def __init__(self):
        super(RnnWith2HiddenSizesModel, self).__init__()
        self.rnn = nn.LSTM(input_size=10, hidden_size=20, num_layers=2)
        self.rnn_two = nn.LSTM(input_size=20, hidden_size=2)

    def forward(self, inp, hc):
        output, _ = self.rnn(inp, hc)
        output2, _ = self.rnn_two(output)
        return output2


inp = torch.randn(5, 3, 10)
h0 = torch.randn(2, 3, 20)
c0 = torch.randn(2, 3, 20)

rnn = RnnWith2HiddenSizesModel()

output = RnnWith2HiddenSizesModel()(inp, (h0, c0))


tensor([[[-0.0305,  0.0327],
     [-0.1534, -0.1193],
     [-0.1393,  0.0474]],

    [[-0.0370,  0.0519],
     [-0.2081, -0.0693],
     [-0.1809,  0.0826]],

    [[-0.0561,  0.0731],
     [-0.2307, -0.0229],
     [-0.1780,  0.0901]],

    [[-0.0612,  0.0891],
     [-0.2378,  0.0164],
     [-0.1760,  0.0929]],

    [[-0.0660,  0.1023],
     [-0.2176,  0.0508],
     [-0.1611,  0.1053]]], grad_fn=<CatBackward>)

答案 1 :(得分:0)

尽管@Mikhail Berlinkov的答案可以根据需要进行工作,但它并没有推广到一般情况下(问题中甚至没有要求),为此,我想提出第二种解决方案:

import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import reduce

class RNNModel(nn.Module):
    def __init__(self, *models):
        super(RNNModel, self).__init__()
        self.models = models

    def forward(self, inp):
        return reduce(lambda arg, model: model(arg, None)[0], self.models, inp)

并且可以称为:

rnn = nn.LSTM(input_size=10, hidden_size=20, num_layers=2)
rnn_two = nn.LSTM(input_size=20, hidden_size=2)
inp = torch.randn(5, 3, 10)

rnn_model = RNNModel(rnn, rnn_two)

output = rnn_model(inp)

output.shape等于预期值(即5,3,2