如何在PyTorch的Sequential模型中提取LSTM层最后隐藏状态

方案1：拆分Sequential插入自定义可导出模块 #

不要使用Lambda层，自定义无参数的nn.Module子类实现提取LSTM最后隐藏状态的逻辑，该类的操作完全符合静态图导出要求，可被正常转换：

import torch
import torch.nn as nn

# 自定义无参数模块，专门提取LSTM最后一层的最后时间步隐藏状态，支持ONNX导出
class ExtractLSTMLastHidden(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x):
        # x为LSTM原生输出，格式为 (output, (h_n, c_n))
        _, (hidden, _) = x
        # 取最后一层隐藏态，输出形状为 (batch_size, hidden_size)
        return hidden[-1]

重构后的Sequential模型结构如下，逻辑和原结构完全对齐：

model = torch.nn.Sequential(
        torch.nn.LSTM(40, 256, 3, batch_first=True),
        ExtractLSTMLastHidden(), # 插入自定义模块承接LSTM输出
        torch.nn.Linear(256, 256),
        torch.nn.ReLU()
)

该自定义模块没有训练参数，不会干扰预训练权重加载，如果你的预训练权重是针对原Sequential结构保存的，加载时手动调整state_dict的key序号，将原Linear、ReLU层对应的key序号加1即可，也可以直接按层名赋值权重避免序号不匹配问题。

方案2：走标准转换链路确保兼容性 #

转换链路按 PyTorch -> ONNX -> TensorFlow SavedModel -> TensorFlow.js 走即可，步骤如下：

• 加载完Resemblyzer预训练权重后，将模型设为eval模式，用相同的测试输入验证PyTorch侧的推理结果符合预期
• 导出为ONNX格式，选择11以上的opset版本，开启动态轴支持可变批次、可变序列长度：

dummy_input = torch.randn(1, 100, 40) # 形状对应 (batch_size, seq_len, feature_dim)，可按实际输入调整
torch.onnx.export(
    model,
    dummy_input,
    "resemblyzer.onnx",
    opset_version=13,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size", 1: "seq_len"}, "output": {0: "batch_size"}}
)

• 用onnx-tf工具将ONNX模型转换为TensorFlow SavedModel格式，确保onnx、onnx-tf、tensorflow版本匹配
• 用TensorFlow.js官方的tensorflowjs_converter工具将SavedModel转换为TF.js支持的格式，即可直接在JS环境加载使用

注意事项 #

• 整个模型的forward逻辑不要加入依赖张量值的动态控制流（比如根据张量值判断的if分支、长度随张量变化的for循环），上述自定义模块的操作都是静态图原生支持的操作，不会触发转换错误
• 转换完成后建议在TF.js环境用和PyTorch侧相同的输入做结果校验，误差控制在1e-5以内即为正常

内容的提问来源于stack exchange，提问作者Cooper