LSTM层后接BatchNorm1d层出现维度不匹配错误求助

问题背景 #

模型架构 #

self.lstm1 = nn.LSTM(in_channels, hidden_channels, num_layers, batch_first=True, dropout=dropout_prob)
self.batchnorm1 = nn.BatchNorm1d(hidden_channels)

self.lstm2 = nn.LSTM(hidden_channels, hidden_channels, num_layers, batch_first=True, dropout=dropout_prob)
self.batchnorm2 = nn.BatchNorm1d(hidden_channels)

self.lstm3 = nn.LSTM(hidden_channels, hidden_channels, num_layers, batch_first=True, dropout=dropout_prob)
self.batchnorm3 = nn.BatchNorm1d(hidden_channels)

self.fc1 = nn.Linear(hidden_channels, out_channels)

参数与维度信息 #

in_channels = 105
hidden_channels = 128
batch_size = 32
# 输入数据维度：(32,770,105) → (batch_size, seq_len, feature_dim)
# LSTM1输出维度：(32,770,128) → (batch_size, seq_len, hidden_dim)

运行错误 #

RuntimeError: running_mean should contain 770 elements not 128

用户曾尝试将LSTM输出permute为(32,128,770)，但仍出现其他错误。

错误原因分析 #

PyTorch的nn.BatchNorm1d默认对第二个维度做归一化，要求输入格式为(batch_size, num_features, seq_len)或(batch_size, num_features)。

而你的LSTM输出是(batch_size, seq_len, hidden_dim)，直接传入BatchNorm1d(hidden_channels)时，BatchNorm会把第二个维度（770，即序列长度）误认为是num_features，因此要求running_mean的长度为770，但初始化时指定的是128，最终导致维度不匹配错误。

解决方案 #

需要在BatchNorm前后调整维度顺序，让hidden_channels（即num_features）处于第二个维度，处理完成后再调回LSTM需要的输入格式：

修改后的前向传播示例 #

def forward(self, x):
    # 输入x shape: (32,770,105)
    x, _ = self.lstm1(x)  # LSTM输出shape: (32,770,128)
    
    # 调整维度适配BatchNorm1d：(batch_size, hidden_channels, seq_len)
    x = x.permute(0, 2, 1)  # shape变为(32,128,770)
    x = self.batchnorm1(x)
    # 调回LSTM要求的输入格式：(batch_size, seq_len, hidden_channels)
    x = x.permute(0, 2, 1)  # shape变回(32,770,128)
    
    # 后续LSTM层重复相同操作
    x, _ = self.lstm2(x)
    x = x.permute(0, 2, 1)
    x = self.batchnorm2(x)
    x = x.permute(0, 2, 1)
    
    x, _ = self.lstm3(x)
    x = x.permute(0, 2, 1)
    x = self.batchnorm3(x)
    x = x.permute(0, 2, 1)
    
    # 根据任务需求处理全连接层输入：
    # 序列预测任务：直接传入即可，输出为(32,770, out_channels)
    # 序列分类任务：需对时间步维度做处理（如取最后一步）
    # x = x[:, -1, :]  # 取最后一个时间步的特征，shape变为(32,128)
    x = self.fc1(x)
    return x

关于permute后仍报错的说明 #

如果你之前仅做了一次permute就传入下一层LSTM，会因为输入格式不符合LSTM要求（batch_first=True时需要(batch_size, seq_len, feature_dim)）而报错，必须再permute回原维度顺序才能正常输入到下一层LSTM。

另外，全连接层的处理需匹配任务类型：若为序列分类，需先对时间步维度做池化或取特定时刻特征，否则全连接层会对每个时间步单独输出，可能不符合任务预期。

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