python - 将循环层附加到具有不同隐藏大小的 PyTorch LSTM 模型

Question

我正在使用 PyTorch 开发用于序列分析的 BI-LSTM 模型。为此，我正在使用 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但找不到我需要的东西。

有什么建议吗？

Answer 1

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

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>)