如何用PyTorch构建多输出多分类（多任务）神经网络？

核心问题 #

你当前的模型最后一层仅输出3个值，对应每个样本3个输出目标的单值预测，但这不符合每个输出目标对应3个类别概率的需求。要实现(4,3,3)的输出，需要让模型为每个目标输出3个类别logits，再整理维度并对每个目标的类别做Softmax归一化。

具体修正步骤 #

1. 调整最后一层输出维度：3个输出目标 × 每个目标3个类别 = 9个输出节点，让模型先输出形状为(4,9)的张量。
2. 重塑输出张量形状：将(4,9)的张量重新调整为(4,3,3)，对应[样本数, 输出目标数, 类别数]的维度结构。
3. 修正Softmax计算维度：对每个输出目标的类别维度（dim=2）做Softmax，确保每个目标的3个类别概率和为1。
4. 目标数据预处理：原始目标值是(-1,0,1)，需要转换为(0,1,2)的索引格式，适配PyTorch损失函数的输入要求。

修正后的完整代码 #

import pandas as pd
from torch import nn 
import torch

feature_data = {
    'A': [1, 2, 3, 4],
    'B': [5, 6, 7, 8],
    'C': [9, 10, 11, 12],
    'D': [13, 14, 15, 16],
    'E': [17, 18, 19, 20],
    'F': [21, 22, 23, 24]
}

target_data = {
    'Col1': [1, -1, 0, 1],
    'Col2': [-1, 0, 1, -1],
    'Col3': [-1, 0, 1, 1]
}

# Create the DataFrame
train_feature_data = pd.DataFrame(feature_data) 
train_target_data = pd.DataFrame(target_data)
device = "cuda" if torch.cuda.is_available() else "cpu"

# 预处理目标数据：将(-1,0,1)转换为(0,1,2)的索引
train_target_data = train_target_data + 1
y_train = torch.tensor(train_target_data.to_numpy()).type(torch.long).to(device)

# 修正后的模型
class MyModel(nn.Module):
  def __init__(self, inputs=6, l1=12, num_targets=3, num_classes=3):
      super().__init__()
      self.num_targets = num_targets
      self.num_classes = num_classes
      self.sequence = nn.Sequential(
        nn.Linear(inputs, l1),
        nn.Linear(l1, num_targets * num_classes)  # 输出节点数=目标数×类别数
    )
      
  def forward(self, x):
      x = self.sequence(x)
      # 重塑形状为[样本数, 目标数, 类别数]
      x = x.view(-1, self.num_targets, self.num_classes)
      # 对每个目标的类别维度做Softmax
      x = nn.Softmax(dim=2)(x)
      return x
    
x_train = torch.tensor(train_feature_data.to_numpy()).type(torch.float).to(device)
model = MyModel(inputs=6, l1=12, num_targets=3, num_classes=3).to(device)
output = model(x_train)
print("输出形状:", output.shape)  # 输出: torch.Size([4, 3, 3])

额外说明 #

• 训练时若使用nn.CrossEntropyLoss，建议去掉模型内的Softmax层——因为CrossEntropyLoss会自动对输入的logits做Softmax计算，这样能避免数值稳定性问题；如果要保留模型输出概率的功能，可改用nn.NLLLoss配合模型内的Softmax。

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