PyTorch搭建回归神经网络输出层设为1时报错问题咨询

报错本质是任务类型、损失函数、标签处理逻辑三者完全不匹配，这也是输出层维度设为100时不报错、设为1就触发异常的核心原因：

• 你的目标是完成h1_hemoglobin连续值回归预测，但代码中使用的nn.CrossEntropyLoss是多分类任务专属损失函数，该损失强制要求两个条件：
  • 模型输出shape为(batch_size, 类别数)，且类别数必须大于等于2
  • 输入标签为LongTensor类型的整数类别索引，shape为(batch_size,)
• 当输出层维度设为1时，模型输出shape为(batch_size, 1)，既不满足交叉熵损失对类别数的最低要求，同时回归任务的标签是连续数值，并非分类任务的整数类别索引，自然触发维度与类型校验报错。
• 额外逻辑错误：代码中将回归标签转换为torch.LongTensor（整数张量）是分类任务的处理流程，完全不适配回归任务的连续值特性。

仅需针对回归任务特性修改3处核心逻辑即可，模型结构本身的输出层维度设为1是完全正确的：

• 替换损失函数：删除分类用的交叉熵损失，改用回归任务常用的均方误差损失nn.MSELoss()，如果需要降低异常值对训练的影响，也可替换为平均绝对误差损失nn.L1Loss()
• 修正标签张量类型：将标签从整数型LongTensor改为浮点型张量，同时调整shape对齐模型输出，避免维度不匹配
• 修正不规范调用：原代码中手动调用model.forward()的写法不符合PyTorch最佳实践，直接通过model(x)调用即可，框架会自动执行forward方法；回归任务最后一层不需要加softmax、sigmoid等分类场景用的激活函数，保持线性输出即可。

修正后可运行代码 #

数据处理部分 #

from sklearn.model_selection import train_test_split
import torch
import torch.nn as nn
import torch.nn.functional as F

# 划分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)
# 张量转换：特征统一用float32，回归标签转float32并调整维度对齐输出
X_train = torch.tensor(X_train, dtype=torch.float32)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).reshape(-1, 1)
y_test = torch.tensor(y_test, dtype=torch.float32).reshape(-1, 1)

模型与训练部分 #

class ANN_Model(nn.Module):
    def __init__(self, input_features=4, hidden1=20, hidden2=20, out_features=1):
        super().__init__()
        self.f_connected1 = nn.Linear(input_features, hidden1)
        self.f_connected2 = nn.Linear(hidden1, hidden2)
        self.out = nn.Linear(hidden2, out_features)
    def forward(self, x):
        x = F.relu(self.f_connected1(x))
        x = F.relu(self.f_connected2(x))
        x = self.out(x)
        return x

# 实例化模型
model = ANN_Model()
# 回归任务使用MSE损失
loss_function = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

epochs = 500
final_losses = []
for i in range(epochs):
    y_pred = model(X_train)
    loss = loss_function(y_pred, y_train)
    final_losses.append(loss.item())
    if (i+1) % 10 == 0:
        print(f"Epoch number: {i+1} and the loss: {loss.item():.4f}")
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

输出层维度设为100时不报错的原因是，100满足交叉熵损失对“类别数≥2”的强制校验，但此时模型是按照100类分类任务的逻辑训练，输出的是100个类别的分类logit，和预测连续血红蛋白值的回归目标完全无关，因此训练得到的模型没有实际使用价值。
推理阶段模型直接输出的值就是h1_hemoglobin的预测结果，不需要执行argmax、softmax等分类任务专属的后处理操作。

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