PyTorch交叉熵损失权重与类别数不匹配问题求助

核心问题原因 #

你的pred尺寸是[B,1]，PyTorch的CrossEntropyLoss会把这个形状判定为仅1个类别的输出（对应单通道logit），但你的target实际包含0、1、2三类，两者不匹配，导致权重张量shape=[3]和判定的类别数冲突，触发报错。

具体解决步骤 #

1. 调整模型输出维度
   把模型最后一层的输出通道数改为3，让pred的尺寸变为[B,3]——这是CrossEntropyLoss处理3分类任务的标准输入格式，每个样本对应3个类别的logit值。
   示例代码：

   # 假设原模型最后一层是nn.Linear(in_features, 1)，修改为：
   model.fc = nn.Linear(in_features, 3)
   

2. 正确设置类别权重
   当pred为[B,3]时，形状为[3]的权重张量就符合要求了，可根据类别频率设置权重（比如给低频的1、2类更高权重）：

   import torch
   import torch.nn as nn
   
   # 示例：手动设置权重（也可根据数据集统计的类别占比计算）
   weights = torch.tensor([0.2, 0.4, 0.4], dtype=torch.float32)
   criterion = nn.CrossEntropyLoss(weight=weights.to(device))
   

3. 验证输入输出匹配
   确保target的取值范围是0~2（对应3个类别），且pred的第二个维度固定为3，此时损失函数会自动按权重计算每个类别的损失占比。

额外说明 #

如果暂时不想修改模型输出，也可以用BCEWithLogitsLoss配合one-hot编码的target，但这种方式更适合二分类或多标签任务，对于你的单标签3分类场景，优先选择调整模型输出适配CrossEntropyLoss的方案更简洁。

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