自定义HuggingFace Trainer序列损失时遇两类错误求助

问题描述 #

自定义CustomTrainer的compute_loss方法时遇到两类错误：

1. 保留类型转换代码时，抛出错误：

IndexError: Target 17 is out of bounds.

2. 移除类型转换代码后，抛出新错误：

RuntimeError: Expected floating point type for target with class probabilities, got Long

原实现代码：

class CustomTrainer(Trainer):
  def compute_loss(self, model, inputs, return_outputs=False):
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda" if use_cuda else "cpu")
    
    # forward pass
    outputs = model(**inputs)
    logits = outputs.get("logits")
   

    # logits: (batch_size, sequence_length, num_classes), reshape it to (batch_size*sequence_length, num_classes)
    predictions = logits.view(-1, model.config.num_labels)
    # shape: (batch_size*sequence_length)
    predictions = predictions.argmax(dim=-1)

    labels = inputs.get("labels").to(device)
    labels =  labels.view(-1) 

    labels = torch.tensor(labels, dtype=torch.float)
    predictions = torch.tensor(predictions, dtype=torch.float, requires_grad=True)

    loss = loss_fct(predictions, labels)


    return (loss, outputs) if return_outputs else loss

错误根源 #

1. 第一个错误（IndexError）：

   • 对logits执行argmax(dim=-1)将其转换为类别索引，再传入损失函数是错误的——训练阶段需要用原始logits（未归一化的概率分布）计算梯度，argmax仅用于推理阶段。
   • 额外用torch.tensor(labels)重新创建张量，可能导致设备不匹配，且如果损失函数是CrossEntropyLoss，传入argmax后的类别索引会和标签维度/范围不匹配，触发越界错误。

2. 第二个错误（RuntimeError）：

   • 移除类型转换后，predictions是argmax生成的long型张量，若损失函数是BCEWithLogitsLoss（要求输入为浮点型logits和浮点标签），会因类型不匹配报错；若用CrossEntropyLoss，则需要输入浮点型logits而非argmax后的类别索引。

修正方案 #

核心原则 #

训练时直接使用模型输出的原始logits计算损失，argmax仅在推理时使用；根据任务类型选择对应损失函数，保证输入张量的类型、维度完全匹配。

场景1：单标签序列分类（用CrossEntropyLoss） #

适用于每个序列位置仅属于一个类别的任务：

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        # 直接复用模型所在设备，无需手动判断
        device = model.device
        
        # 前向传播获取logits
        outputs = model(**inputs)
        logits = outputs.get("logits")
        
        # 调整张量形状：(batch_size*sequence_length, num_labels) 和 (batch_size*sequence_length,)
        logits = logits.view(-1, model.config.num_labels)
        labels = inputs.get("labels").view(-1).to(device)
        
        # 初始化损失函数（若未全局定义）
        loss_fct = torch.nn.CrossEntropyLoss()
        # CrossEntropyLoss要求logits为float型，labels为long型，模型输出和输入标签默认符合该要求
        loss = loss_fct(logits, labels)
        
        return (loss, outputs) if return_outputs else loss

场景2：多标签/概率序列任务（用BCEWithLogitsLoss） #

适用于每个序列位置可属于多个类别，或标签为概率值的任务：

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        device = model.device
        
        outputs = model(**inputs)
        logits = outputs.get("logits")
        
        # 调整形状：(batch_size*sequence_length, num_labels) 和 (batch_size*sequence_length, num_labels)
        logits = logits.view(-1, model.config.num_labels)
        labels = inputs.get("labels").view(-1, model.config.num_labels).to(device).float()
        
        loss_fct = torch.nn.BCEWithLogitsLoss()
        # BCEWithLogitsLoss要求logits和标签均为float型
        loss = loss_fct(logits, labels)
        
        return (loss, outputs) if return_outputs else loss

关键注意事项 #

• 不要用torch.tensor(labels)重新创建张量，直接用inputs["labels"].to(device)即可，避免设备不匹配和冗余数据复制。
• 绝对不要在训练阶段对logits执行argmax，否则会破坏梯度传播链，导致模型无法学习。
• 确保labels的形状与调整后的logits完全匹配：单标签任务labels为一维张量，多标签任务labels为二维张量（与logits最后一维维度一致）。

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