3D输入下nn.Linear与nn.BatchNorm1d的参数设置疑问

你踩的这个坑特别典型，很多刚接触PyTorch批量归一化的同学都会在这里卡壳。核心原因是**nn.BatchNorm1d对不同维度输入的“特征维度”定义不一样**，咱们一步步拆解明白：

先搞懂nn.BatchNorm1d的维度逻辑 #

PyTorch的nn.BatchNorm1d并不是固定盯着最后一个维度当“特征维度”，而是根据输入的总维度数来判断：

• 当输入是2D张量（形状为(batch_size, num_features)）：比如你提到的(2,70)经过Linear变成(2,20)，这时候num_features就是最后一个维度的大小（20），BN会对每个特征维度，在整个batch的样本上计算均值和方差。
• 当输入是3D张量（形状为(batch_size, num_features, sequence_length)）：这时候num_features是中间的那个维度！BN会把batch_size * sequence_length所有元素当成一个集合，对每个num_features维度单独计算均值和方差。

你的示例为什么一个正常一个报错？ #

看你的两个例子：

• 示例1中，输入经过Linear后是(2,50,20)，PyTorch把它识别为3D输入：batch_size=2，num_features=50，sequence_length=20。所以nn.BatchNorm1d(50)刚好匹配中间的50，运行完全正常。
• 示例2中你设了nn.BatchNorm1d(20)，但PyTorch预期输入的num_features是50（中间维度），所以会抛出RuntimeError: running_mean should contain 50 elements not 20——因为BN层维护的running_mean长度等于num_features，和输入的中间维度不匹配。

如果想对Linear输出的20个特征做BN怎么办？ #

如果你本来的需求是：对每个位置的20个输出特征做归一化（也就是把(2,50,20)里的20当成特征维度），那需要先调整输入形状，把它变成2D张量再用BN，之后恢复形状即可。比如：

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.bn11 = nn.BatchNorm1d(20)
        self.fc11 = nn.Linear(70,20)
    def forward(self, inputs):
        out = self.fc11(inputs)  # 形状: (2,50,20)
        # 合并前两个维度，转成2D张量
        out = out.flatten(0, 1)  # 形状: (2*50=100, 20)
        out = torch.relu(self.bn11(out))
        # 恢复原有的3D形状
        out = out.unflatten(0, (2, 50))  # 形状: (2,50,20)
        return out

这样nn.BatchNorm1d(20)就会对每个特征维度（20个），在100个“样本”（原batch的每个元素+每个位置）上做归一化，完全符合你的预期。

总结一下 #

• 2D输入用nn.BatchNorm1d(num_features)：num_features对应最后一个维度。
• 3D输入用nn.BatchNorm1d(num_features)：num_features对应中间的维度。
• 如果要对3D输入的最后一个维度做BN，需要先把前两个维度合并成一个batch维度，处理后再恢复形状。

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