PyTorch中样本类别数可变的分类任务交叉熵损失处理方法问询

针对分类任务中每个样本对应类别数量不同的场景，PyTorch没有直接支持批量输入的统一API，但有标准且高效的处理方式，核心逻辑是逐个处理样本后聚合损失。

关键前提 #

PyTorch的F.cross_entropy或nn.CrossEntropyLoss有两个核心输入要求：

• logits需为[N, C]形状的张量（N是批次大小，C是固定类别数）
• 标签需传入类索引（而非one-hot向量，若现有标签是one-hot格式，需先转换）

针对示例的处理步骤 #

假设你的示例数据如下：

import torch
import torch.nn.functional as F

# 示例输入
logits_list = [
    torch.tensor([0.2, 0.2, 0.6]),    # 对应3类
    torch.tensor([0.4, 0.1, 0.1, 0.4]), # 对应4类
    torch.tensor([0.2, 0.8])         # 对应2类
]
labels_list = [
    torch.tensor([0, 0, 1]),  # one-hot格式标签
    torch.tensor([1, 0, 0, 0]),
    torch.tensor([1, 0])
]

步骤1：转换标签格式

将one-hot标签转为类索引，适配损失函数的输入要求：

labels_idx = [torch.argmax(label) for label in labels_list]
# 转换结果为：tensor(2), tensor(0), tensor(0)

步骤2：逐个计算损失并聚合

由于每个样本的类别数不同，无法直接堆叠成统一形状的批量张量，因此逐个计算单样本损失后，求批次平均损失：

total_loss = 0.0
batch_size = len(logits_list)

for logits, label in zip(logits_list, labels_idx):
    # 为单样本添加batch维度，匹配F.cross_entropy的输入格式
    loss = F.cross_entropy(logits.unsqueeze(0), label.unsqueeze(0))
    total_loss += loss.item()

# 计算批次平均损失
avg_loss = total_loss / batch_size
print(avg_loss)

效率说明 #

这种逐个计算的方式在PyTorch动态图机制下不会有明显性能损耗，梯度会自动完成累积。如果模型本身输出的就是变长logits（比如动态序列分类场景），可以直接在模型前向传播阶段按样本计算损失，避免额外的列表拼接操作。

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