PyTorch 1.2.0 CrossEntropyLoss报错：仅支持3D空间目标张量，输入维度不符

嘿，我来帮你搞定这个问题！你遇到的报错核心不是target的维度不对，而是模型输出形状不符合MNIST分类任务的要求，导致CrossEntropyLoss误判成了空间分割任务的场景。

让我拆解一下关键逻辑：

• PyTorch的CrossEntropyLoss有两种典型使用场景：
  • 标准分类任务：输入output形状为(batch_size, num_classes)，目标target是一维张量(batch_size,)，每个元素对应样本的类别索引（MNIST里就是0-9的整数）
  • 空间分割任务：输入output形状为(batch_size, num_classes, H, W)，这时候Loss会认为你要对每个像素分类，要求target是三维张量(batch_size, H, W)，每个位置对应该像素的类别

你的情况是：模型处理(32,1,28,28)的输入后，输出的output形状应该是(32,10)（MNIST是10类），但实际你的模型输出可能还是带空间维度的张量（比如(32,10,28,28)）。这时候Loss自动切换到空间分类逻辑，要求target是3维，但你给的target.unsqueeze(1)是(32,1)的2维张量，自然就触发报错了。

1. 恢复target的正确形状 #

首先把错误的target修改代码改回来，MNIST分类任务的target本身就是(batch_size,)的一维张量，完全符合标准分类要求，不需要额外加维度：

loss = loss_fn(output, target)  # 去掉target.unsqueeze(1)

2. 修改模型，确保输出为(batch_size, num_classes) #

你的模型需要在最后把空间维度的特征图展平成一维，再通过全连接层输出10个类别的logits。分两种情况调整：

如果是带卷积的模型：

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(1, 32, kernel_size=3)
        # 计算展平后的特征数：32 channels * 26*26 空间维度
        self.fc = nn.Linear(32*26*26, 10)
    
    def forward(self, x):
        x = self.conv(x)
        # 从第1维开始展平（第0维是batch）
        x = x.flatten(1)  # 等价于torch.flatten(x, start_dim=1)
        x = self.fc(x)  # 现在x的形状是(32,10)，符合要求
        return x

如果是纯MLP模型：

需要先把输入的(batch_size,1,28,28)展平成(batch_size, 784)，再送入全连接层：

class MyMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 256)
        self.fc2 = nn.Linear(256, 10)
    
    def forward(self, x):
        x = x.flatten(1)  # 把(32,1,28,28)转为(32,784)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)  # 输出(32,10)
        return x

3. 验证形状是否正确 #

你可以在train函数里加两行打印，确认输出和目标的形状是否符合预期：

output = model(inputs)
print("Output shape:", output.shape)
print("Target shape:", target.shape)

正常情况下，output应该是(32,10)，target是(32,)，这样Loss就能正常计算了。

内容的提问来源于stack exchange，提问作者J.C