PyTorch中LSTM连接Linear层实现MFCC分类的维度报错问题

LSTM输出维度规则（batch_first=True配置下） #

首先明确PyTorch中单向、非双向LSTM的返回值维度，所有报错均由维度不匹配导致：

• out：所有时间步的隐层输出，维度为(batch_size, seq_len, hidden_size)，对应你的输入配置，维度为(16, 60, 128)
• ht：最后一个时间步的隐层状态，维度为(num_layers * num_directions, batch_size, hidden_size)，你当前配置是1层单向LSTM，所以维度为(1, 16, 128)
• ct：最后一个时间步的细胞状态，维度和ht一致，整序列分类任务一般不需要用到。

三类报错的根因 #

• 报错1（ht.contiguous().view(16,-1)触发）：硬编码batch维度做reshape时没有先调整ht的维度顺序，同时硬编码batch=16的写法无法适配不同batch size的场景，最终传入全连接层的特征维度和权重维度不匹配。
• 报错2（out.contiguous().view(16,-1)触发）：把所有60个时间步的特征全部展平，得到的特征维度是16*(60*128) = 16*7680，但全连接层定义的输入维度是128，二者维度完全无法对齐。
• 报错3（不做维度处理直接传入全连接触发）：全连接层默认对输入张量的最后一维做线性变换，直接传入形状为(16,60,128)的out，经过全连接后输出维度为(16,60,32)，而nn.CrossEntropyLoss要求分类logits的形状为(batch_size, class_num)即(16,32)，标签形状为(16,)，维度不匹配直接报错。

正确实现代码 #

整序列分类任务只需要取最后一个时间步的隐层特征喂入全连接层即可，不需要展平所有时间步的输出，修正后的模型代码如下：

import torch
import torch.nn as nn

class model(nn.Module):
  def __init__(self,ninp,num_layers,class_num,nhid=128):
      super().__init__()
      
      self.lstm_nets = nn.LSTM(input_size=ninp,hidden_size=nhid,num_layers=num_layers,
      batch_first=True,dropout=0.2,bidirectional=False)
      self.FC = nn.Linear(nhid,class_num)
      self.tanh = nn.Tanh()
      # 注意：nn.CrossEntropyLoss内部已集成LogSoftmax+NLLLoss，不需要额外加Softmax/LogSoftmax层，否则会导致损失计算异常
      
  def forward(self,X):
      out, (ht, ct) = self.lstm_nets(X)
      # 压缩num_layers维度，得到形状为(batch_size, hidden_size)的最后一步隐特征
      # 不要硬编码batch size，用squeeze自动适配任意batch大小
      last_hidden = ht.squeeze(0)
      out = self.tanh(last_hidden)
      logits = self.FC(out)
      return logits

model = model(ninp=40,num_layers=1,class_num=32,nhid=128)
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.5e-4)

可选实现方式 #

如果你希望融合所有时间步的特征做分类，也可以用全局平均池化替代取最后一步隐状态的操作，这种方式可以适配任意长度的输入序列：

def forward(self,X):
    out, (ht, ct) = self.lstm_nets(X)
    # 对时间步维度做平均池化，输出形状为(batch_size, hidden_size)
    out = torch.mean(out, dim=1)
    out = self.tanh(out)
    logits = self.FC(out)
    return logits

注意事项 #

• 禁止在reshape操作中硬编码batch size数值（比如.view(16,-1)），否则验证、推理阶段batch size变化时会直接触发维度报错。
• 使用nn.CrossEntropyLoss时不要在网络末尾添加Softmax/LogSoftmax层，避免损失计算重复导致模型收敛异常。
• 只有逐帧标注类的任务（比如语音识别、实体识别）才需要对每个时间步的输出做分类，整序列分类任务只需要用聚合后的时序特征即可。

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