多输出多分类任务中CrossEntropyLoss维度不匹配问题求助

嗨，我来帮你捋捋这个维度不匹配的问题～ 你遇到的RuntimeError: Expected target size [64, 5], got [64, 4]，核心原因是CrossEntropyLoss的输入格式和你的模型输出、标签形状不匹配。

先拆解你的任务：你有4个独立的5分类任务（每个动作对应5个类别），模型需要为每个样本的4个动作各自输出5个类别的概率。那我们得调整输出和损失计算的方式，让它符合PyTorch损失函数的要求：

问题根源分析 #

PyTorch的CrossEntropyLoss默认期望的输入格式是：

• 预测输出：[N, C]，其中N是批量大小，C是类别数
• 标签：[N]，每个元素是对应样本的类别索引

但你现在把输出reshape成了[64,4,5]，标签是[64,4]，直接喂给损失函数的话，它会误以为你的类别数是4，而标签的维度不对，所以报错。

具体解决步骤 #

1. 去掉模型里的Softmax

CrossEntropyLoss内部已经整合了LogSoftmax和NLLLoss，提前加Softmax不仅会浪费计算，还可能导致数值不稳定，所以把模型forward里的torch.softmax去掉：

def forward(self, x):
    out = self.layer1(x)
    out = self.relu(out)
    out = self.layer2(out)
    out = self.relu(out)
    out = self.output_layer(out)
    # 这里直接返回原始输出，不用做softmax
    return out

2. 调整训练循环里的输出与标签维度

我们需要把4个分类任务“扁平化”成一个大的批量任务：

• 将模型输出从[batch_size, num_actions*5] reshape为[batch_size*num_actions, 5]
• 将标签从[batch_size, num_actions] flatten为[batch_size*num_actions]

修改训练循环里的损失计算部分：

for inputs, labels in train_loader:
    # Forward pass
    outputs = model(inputs)
    # 调整输出形状：[batch_size, num_actions*5] → [batch_size*num_actions, 5]
    outputs = outputs.view(-1, 5)
    # 把标签扁平化：[batch_size, num_actions] → [batch_size*num_actions]
    labels = labels.flatten()
    
    loss = criterion(outputs, labels)
    
    # 后面的反向传播和优化不变
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    total_loss += loss.item()

3. 检查标签映射的小问题

你代码里的注释写的是If value is -2, map it to 1，但实际代码里是np.where(Y == -2, 0, ...)，这里注释和代码逻辑不一致，记得确认下你的类别映射是否符合预期哦～

验证修改后的维度 #

修改后，对于batch_size=64的情况：

• 输出形状：64*4=256个样本，每个对应5个类别，即[256,5]
• 标签形状：[256]，正好匹配CrossEntropyLoss的要求，维度不匹配的问题就解决啦！

备注：内容来源于stack exchange，提问作者May