PyTorch训练MNIST识别CNN时遇张量尺寸不匹配报错求助

嘿，我帮你分析下这个问题哈！你遇到的警告和错误本质上是同一个原因：模型输出和标签的维度/格式不匹配。

问题根源 #

你的模型最后一层输出的是形状为 torch.Size([64, 10]) 的张量（对应每个样本的10个类别预测值），但MNIST数据集加载的标签是形状为 torch.Size([64]) 的整数索引（比如0-9的数字）。MSE Loss要求输入和目标的形状完全一致，直接把这两个张量丢进去计算，PyTorch会尝试广播匹配，结果就触发了警告和后续的维度不匹配错误。

两种解决方案 #

方案1：坚持使用MSE Loss，将标签转为One-Hot编码

既然MSE需要和模型输出同维度的目标，你可以把整数标签转换成One-Hot编码（每个标签变成长度为10的向量，对应类别位置为1，其余为0），同时要转成浮点类型（和模型输出的类型一致）：

修改训练代码中的loss计算部分：

import torch.nn.functional as F

# ... 其他训练代码不变 ...
for images, labels in trainloader:
    optimizer.zero_grad()
    output = net(images)
    # 将标签转为One-Hot编码，并转换为float类型
    labels_onehot = F.one_hot(labels, num_classes=10).float()
    loss = criterion(output, labels_onehot)
    loss.backward()
    optimizer.step()
    running_loss += loss.item()
# ... 其他代码不变 ...

方案2：更推荐——改用分类任务专用的交叉熵损失

MNIST是分类任务，MSE Loss其实不是最优选择，**交叉熵损失（CrossEntropyLoss）**才是分类任务的标准损失函数，它可以直接处理整数类型的标签，不需要手动转One-Hot，而且训练效果更好：

只需要替换损失函数的定义即可，训练代码的其他部分完全不用改：

# 替换原来的MSE Loss为交叉熵损失
criterion = nn.CrossEntropyLoss()

补充说明：CrossEntropyLoss内部已经包含了Softmax层，你的模型最后一层直接输出logit值就可以，不需要额外加Softmax激活，完美适配这个损失函数。

修改后的完整训练代码（方案2示例） #

net = Net()
print(net)
# 使用交叉熵损失
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001)
epochs = 3
for epoch in range(epochs):
    running_loss = 0.0
    for images, labels in trainloader:
        optimizer.zero_grad()
        output = net(images)
        loss = criterion(output, labels)  # 直接用原始整数标签
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    else:
        print(f"Training loss: {running_loss/len(trainloader)}")
print('Finished training')

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