仅从 pytorch 顺序中的 lstm 层检索最后一个隐藏状态答案

【问题标题】：Retrieve only the last hidden state from lstm layer in pytorch sequential仅从 pytorch 顺序中的 lstm 层检索最后一个隐藏状态
【发布时间】：2021-11-25 08:21:42
【问题描述】：

我有一个 pytorch 模型：

model = torch.nn.Sequential(
        torch.nn.LSTM(40, 256, 3, batch_first=True),
        torch.nn.Linear(256, 256),
        torch.nn.ReLU()
)

对于 LSTM 层，我只想从批次中检索最后一个隐藏状态以通过其余层。例如：

_, (hidden, _) = lstm(data)
hidden = hidden[-1]

不过，该示例仅适用于子类模型。我需要以某种方式在 nn.Sequential() 模型上执行此操作，这样当我保存它时，它可以正确地转换为 tensorflow.js 模型。我无法在 tensorflow.js 中制作和训练这个模型的原因是因为我正在尝试在 tensorflow.js 中实现这个 repo：Resemblyzer，同时仍然使用与在 pytorch 中制作的预训练 Resemblyzer 模型相同的权重一个子类模型。我曾想过使用 torchvisions.transformations.Lambda() 转换，但我认为这会使其与 tensorflow.js 不兼容。有什么方法可以在保证模型正确转换的同时实现这一点？

【问题讨论】：

标签： pytorch lstm tensorflow.js tensorflowjs-converter

【解决方案1】：

您可以拆分顺序，但只能在推理模型的前向定义中这样做。一旦定义：

model = nn.Sequential(nn.LSTM(40, 256, 3, batch_first=True),
                      nn.Linear(256, 256),
                      nn.ReLU())

你可以拆分它：

>>> lstm, fc = model[0], model[1:]

然后分两步推断：

>>> out, (hidden, _) = lstm(data)
>>> hidden = hidden[-1]
>>> out = fc(out)       # <- or fc(out[-1]) depending on what you want

【讨论】：