如何修改PyTorch MNIST分类CNN模型从输出层获取连续值并解决MSE报错

报错原因 #

你当前报错的核心原因是输出张量和标签张量维度不匹配：原分类模型输出维度为[batch_size, 10]对应10个类别的概率分布，而回归任务的标签维度为[batch_size]，计算MSELoss时无法完成维度对齐。

需要修改的内容 #

1. 修改模型输出层

将全连接层的输出维度从10改为1，每个样本仅输出1个连续值：

# 原代码
# self.out = nn.Linear(32 * 7 * 7, 10)
# 修改后
self.out = nn.Linear(32 * 7 * 7, 1)

2. 调整标签维度与类型

计算损失前，将标签从[batch_size]调整为[batch_size, 1]，同时转为浮点型适配MSELoss的计算要求：

# 计算损失时修改为以下代码
loss = criterion(output, b_y.float().unsqueeze(1))

3. 可选优化

如果希望输出值限制在0~9的合理范围内，可在forward的输出部分添加截断逻辑：

def forward(self, x):
    x = self.conv1(x)
    x = self.conv2(x)
    x = x.view(x.size(0), -1)       
    output = self.out(x)
    # 新增截断逻辑，限制输出范围为0-9
    output = torch.clamp(output, min=0, max=9)
    return output, x

修改完成后模型即可输出符合要求的连续值，完成回归任务的训练。

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