如何在BERT损失函数中添加权重以提升NER任务F1值？

针对你遇到的不平衡数据集下NER模型F1分数瓶颈的问题，通过给损失函数添加类别权重是直接有效的解决方案，以下是具体实现步骤：

1. 计算类别权重 #

类别权重的核心是让模型更关注占比低的类别，常用两种计算方式：

• 手动计算：根据你给出的类别占比，取占比的倒数作为权重：
  class1权重≈1/0.554≈1.81；class2权重≈1/0.271≈3.69；class3权重≈1/0.185≈5.41
• 自动计算（推荐）：用sklearn工具从训练数据中提取标签后自动生成平衡权重：

from sklearn.utils.class_weight import compute_class_weight
import numpy as np
import torch

# 从训练数据集提取所有标签（展平序列标注的标签）
all_labels = np.concatenate([example["labels"] for example in train_dataset])
# 计算平衡权重
class_weights = compute_class_weight(
    class_weight='balanced',
    classes=np.unique(all_labels),
    y=all_labels
)
# 转为PyTorch张量
class_weights = torch.tensor(class_weights, dtype=torch.float)

注意：权重数组的顺序必须和你的标签字典索引完全对应（比如标签字典{"class1":0, "class2":1, "class3":2}，权重数组第一个元素对应class1）。

2. 自定义带权重损失的Trainer #

Transformers默认Trainer使用无权重的交叉熵损失，需要自定义compute_loss方法注入类别权重：

from transformers import Trainer
import torch.nn.functional as F

class WeightedNER_Trainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        # 取出标签，分离输入
        labels = inputs.pop("labels")
        # 前向传播得到模型输出
        outputs = model(**inputs)
        logits = outputs.logits
        
        # 序列标注任务中，logits形状是(batch_size, seq_len, num_labels)，labels是(batch_size, seq_len)
        # 展平后计算带权重的交叉熵损失
        loss = F.cross_entropy(
            logits.view(-1, self.model.config.num_labels),
            labels.view(-1),
            weight=class_weights.to(logits.device)  # 确保权重和模型在同一设备
        )
        
        return (loss, outputs) if return_outputs else loss

3. 使用自定义Trainer训练 #

用你自定义的WeightedNER_Trainer代替原Trainer即可：

# 初始化自定义训练器（参数和你原来的Trainer一致）
trainer = WeightedNER_Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer
)

# 开始训练
trainer.train()

额外注意事项 #

• 如果你的NER任务包含背景类（O），可以根据其占比决定是否加权重，通常背景类占比高的话，权重设为1即可。
• 除了损失权重，还可以配合过采样少数类样本、调整学习率等策略进一步优化F1分数。

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