使用HuggingFace构建QA Pipeline遇KeyError: 'eval_loss'，求保存最佳模型方法

问题背景 #

使用HuggingFace框架构建问答（QA）Pipeline时，训练阶段出现KeyError: 'eval_loss'错误，配置的TrainingArguments如下：

args = TrainingArguments(f'model_training',
                      evaluation_strategy="epoch",
                      label_names = ["start_positions", "end_positions"],
                      logging_steps = 1,
                      learning_rate=2e-5,
                      num_train_epochs=epochs,
                      save_total_limit = 2,
                      load_best_model_at_end=True,
                      save_strategy="epoch",
                      logging_strategy="epoch",
                      report_to="none",
                      weight_decay=0.01,
                      fp16=True,
                      push_to_hub=False)

错误栈信息：

Traceback (most recent call last):
  File "qa_pipe.py", line 286, in <module>
    pipe.training(train_d, val_d, epochs = 2)
  File "qa_pipe.py", line 263, in training
    self.trainer.train()
  File "/home/admin/qa/lib/python3.7/site-packages/transformers/trainer.py", line 1505, in train
    ignore_keys_for_eval=ignore_keys_for_eval,
  File "/home/admin/qa/lib/python3.7/site-packages/transformers/trainer.py", line 1838, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/home/admin/qa/lib/python3.7/site-packages/transformers/trainer.py", line 2090, in _maybe_log_save_evaluate
    self._save_checkpoint(model, trial, metrics=metrics)
  File "/home/admin/qa/lib/python3.7/site-packages/transformers/trainer.py", line 2193, in _save_checkpoint
    metric_value = metrics[metric_to_check]
KeyError: 'eval_loss'

错误原因 #

当设置load_best_model_at_end=True时，Trainer默认以eval_loss作为筛选最佳模型的指标，但如果验证阶段未生成eval_loss（比如自定义compute_metrics函数未返回该值、验证数据集字段缺失导致模型无法计算损失），就会触发该KeyError。

解决方法 #

1. 明确指定评估指标并确保返回eval_loss #

在TrainingArguments中显式声明最佳模型的判断指标，同时确保compute_metrics函数返回包含eval_loss的结果字典：

args = TrainingArguments(
    f'model_training',
    evaluation_strategy="epoch",
    label_names = ["start_positions", "end_positions"],
    logging_steps = 1,
    learning_rate=2e-5,
    num_train_epochs=epochs,
    save_total_limit = 2,
    load_best_model_at_end=True,
    save_strategy="epoch",
    logging_strategy="epoch",
    report_to="none",
    weight_decay=0.01,
    fp16=True,
    push_to_hub=False,
    # 新增以下配置
    metric_for_best_model="eval_loss",
    greater_is_better=False  # loss值越小模型性能越好
)

2. 检查验证数据集完整性 #

确认验证数据集val_d包含模型所需的全部输入特征：input_ids、attention_mask、start_positions、end_positions，避免因字段缺失导致模型无法计算验证损失。

3. 升级Transformers库版本 #

旧版本Transformers可能存在eval_loss未被正确记录的Bug，执行命令升级依赖：

pip install --upgrade transformers datasets

4. 临时规避方案（可选） #

如果暂时不需要自动加载最佳模型的功能，可先将load_best_model_at_end设为False，训练完成后手动筛选最佳模型：

args = TrainingArguments(
    # 其他参数不变
    load_best_model_at_end=False
)

最佳模型保存与加载 #

错误修复后，Trainer会自动将最佳模型保存至model_training目录。加载模型时直接读取该目录即可：

from transformers import AutoModelForQuestionAnswering, AutoTokenizer

model = AutoModelForQuestionAnswering.from_pretrained("model_training")
tokenizer = AutoTokenizer.from_pretrained("model_training")

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