如何将LSTM编码器输出传入torch.nn.Linear全连接层训练分类器

你当前的LSTM最终隐藏状态形状为[num_layers, batch_size, hidden_size]即[2, 1, 1024]，要适配torch.nn.Linear要求的(N, *, Hin)输入格式（N为batch size，Hin为输入特征维度），核心是把batch size移到第一维，合理处理层数维度即可，常用的可落地方案有三类：

• 方案1：拼接多层隐藏态（保留全部层特征）
  将两层的隐藏特征在特征维度拼接，不会丢失多层编码信息，适合模型容量充足的场景。
  代码示例：

# 输入h_n形状为[2, 1, 1024]
h_n = h_n.permute(1, 0, 2) # 维度重排为 [batch_size, num_layers, hidden_size] → [1, 2, 1024]
h_n = h_n.flatten(start_dim=1) # 压平层数与隐藏层维度 → [1, 2048]
# 后续可直接传入in_features=2048的Linear层

• 方案2：仅提取最后一层隐藏态（轻量化方案）
  如果只需要最顶层LSTM的编码特征，直接截取最后一层输出即可，特征维度保持1024不变。
  代码示例：

# 取最后一层隐藏态，形状为 [batch_size, hidden_size] → [1, 1024]
last_hidden = h_n[-1, :, :]
# 后续可直接传入in_features=1024的Linear层

• 方案3：多层特征池化融合（平衡特征量与信息完整性）
  对多层隐藏态做均值或最大池化，既融合了多层的特征信息，又不会提升特征维度。
  代码示例：

h_n = h_n.permute(1, 0, 2) # 维度重排为 [1, 2, 1024]
# 均值池化
pooled_h = torch.mean(h_n, dim=1) # 输出形状为 [1, 1024]
# 也可替换为最大池化：pooled_h = torch.max(h_n, dim=1)[0]
# 后续可直接传入in_features=1024的Linear层

以上三类方案对任意batch size都通用，不需要额外调整维度索引。

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