PyTorch含多BatchNorm1d层模型推理报错：期望2D/3D输入，得到1D输入

我用PyTorch构建了带多个BatchNorm1d层的Discriminator模型，代码如下：

import torch.nn as nn

class Discriminator(nn.Module):
    def __init__(self, sequenceLength):
        super(Discriminator,self).__init__()
        self.batchnorm1 = nn.BatchNorm1d(sequenceLength)
        self.batchnorm2 = nn.BatchNorm1d(2*sequenceLength)
        self.linear1 = nn.Linear(sequenceLength, 2*sequenceLength)
        self.conv2 = nn.Conv1d(1, 1,kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv1d(1, 1,kernel_size=3, stride=1, padding=1)
        self.linear4 = nn.Linear(2*sequenceLength, 1)
        self.relu = nn.ReLU(0.01)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        out = self.batchnorm1(x)
        out = self.linear1(out)
        out = self.relu(out)
        out = self.batchnorm2(out)
        out = out.unsqueeze(1)
        out = self.conv2(out)
        out = self.sigmoid(out)
        out = self.conv3(out)
        out = self.relu(out)
        out = out.squeeze()
        out = self.batchnorm2(out)
        out = self.linear4(out)
        out = self.sigmoid(out)
        return out

推理代码如下：

import torch

Discriminator = torch.load('disc.pth', map_location=torch.device('cpu'))
Discriminator.eval()
embededSeq = Embedding.EmbedOne('sample data')
embededSeq = torch.tensor(embededSeq).float()
embededSeq = embededSeq.unsqueeze(0)
score = PosDiscriminator(embededSeq).detach().numpy()[0]

推理时在out = self.batchnorm2(out)行触发错误：expected 2D or 3D input (got 1D input)。怀疑是前一行out = out.squeeze()导致，但训练时代码运行正常，仅推理阶段出现问题，恳请帮忙排查错误原因。

核心原因 #

训练时输入的是批量数据（例如shape为(batch_size, seq_len)），经过卷积后的张量shape为(batch_size, 1, 2*seq_len)，执行无参数squeeze()后会去掉维度为1的轴，得到(batch_size, 2*seq_len)的2D张量，符合BatchNorm1d对输入维度的要求。

但推理时输入的是单样本：embededSeq.unsqueeze(0)后shape为(1, seq_len)，经过卷积后张量shape变为(1,1,2*seq_len)，此时无参数squeeze()会把所有维度为1的轴全部移除，最终得到(2*seq_len,)的1D张量，而BatchNorm1d要求输入必须是2D（(N, C)）或3D（(N, C, L)）格式，因此触发错误。

解决方案 #

修改forward函数中的squeeze()调用，指定仅移除卷积添加的通道维度（第1维）：

# 替换原代码中的 out = out.squeeze()
out = out.squeeze(1)

这样无论输入是单样本还是批量数据，只会去掉索引为1的维度，得到的张量shape为(N, 2*seq_len)，完全符合BatchNorm1d的输入要求。

另外注意推理代码中的变量名问题：加载的模型变量名为Discriminator，但调用时使用了PosDiscriminator，需要统一变量名，避免出现"未定义对象"的错误。

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