如何获取transformers.Trainer训练阶段的epoch级指标？

目前使用transformers.Trainer时，评估阶段的指标可以通过内置方法轻松获取，但训练过程中每个epoch的loss、准确率、召回率、精确率、F1值这类核心指标却没有直接的获取途径，零散的讨论也未给出完整可落地的方案。以下是一套可行的实现方法：

实现思路 #

通过自定义TrainerCallback拦截训练的epoch_end事件，在每个epoch结束后，遍历训练集计算并记录各项指标。该方法无需修改Trainer核心逻辑，仅通过回调机制扩展功能。

具体代码实现 #

from transformers import TrainerCallback, TrainingArguments, Trainer
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
import torch

class TrainMetricsCallback(TrainerCallback):
    def __init__(self, train_dataset):
        self.train_dataset = train_dataset
        self.train_metrics = []  # 存储所有epoch的训练指标

    def on_epoch_end(self, args, state, control, model, **kwargs):
        # 切换模型到评估模式，关闭梯度计算
        model.eval()
        total_loss = 0.0
        all_predictions = []
        all_labels = []

        # 获取训练集的数据加载器
        train_dataloader = Trainer.get_eval_dataloader(args, self.train_dataset)
        
        with torch.no_grad():
            for batch in train_dataloader:
                # 把数据移到模型所在设备
                inputs = {k: v.to(args.device) for k, v in batch.items() if k != "labels"}
                labels = batch["labels"].to(args.device)
                
                # 前向传播计算loss和预测结果
                outputs = model(**inputs)
                total_loss += outputs.loss.item() * labels.size(0)
                
                predictions = torch.argmax(outputs.logits, dim=-1)
                all_predictions.extend(predictions.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())

        # 计算各项指标
        avg_loss = total_loss / len(self.train_dataset)
        accuracy = accuracy_score(all_labels, all_predictions)
        precision, recall, f1, _ = precision_recall_fscore_support(
            all_labels, all_predictions, average="weighted"
        )

        # 整理当前epoch的指标
        epoch_metrics = {
            "epoch": round(state.epoch, 2),
            "train_loss": round(avg_loss, 4),
            "train_accuracy": round(accuracy, 4),
            "train_precision": round(precision, 4),
            "train_recall": round(recall, 4),
            "train_f1": round(f1, 4)
        }
        self.train_metrics.append(epoch_metrics)
        
        # 打印当前epoch的训练指标
        print(f"\nEpoch {epoch_metrics['epoch']} Training Metrics:")
        for metric_name, value in epoch_metrics.items():
            print(f"{metric_name}: {value}")
        
        # 切换回训练模式
        model.train()

# 使用示例
def run_training():
    # 假设已完成model、train_dataset、eval_dataset的初始化
    training_args = TrainingArguments(
        output_dir="./training_results",
        per_device_train_batch_size=8,
        num_train_epochs=5,
        logging_dir="./logs",
        logging_steps=100,
        fp16=True  # 按需开启混合精度训练
    )

    # 初始化自定义回调
    train_metric_callback = TrainMetricsCallback(train_dataset)

    # 初始化Trainer并传入回调
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        callbacks=[train_metric_callback]
    )

    # 启动训练
    trainer.train()

    # 训练结束后可查看所有epoch的指标
    print("\nAll Training Epoch Metrics:")
    for metrics in train_metric_callback.train_metrics:
        print(metrics)

if __name__ == "__main__":
    run_training()

关键说明与优化建议 #

• 指标计算逻辑：代码中使用加权平均的精确率、召回率、F1值，若需二分类场景，可将average参数改为binary；多分类场景可改为macro或micro。
• 性能优化：若训练集规模极大，全量计算指标会增加训练耗时，可抽取训练集的固定比例子集（如10%）来计算近似指标。
• 分布式训练适配：如果使用多GPU分布式训练，需通过torch.distributed.all_gather聚合各进程的预测结果与loss，避免指标计算偏差。
• 日志持久化：可将epoch_metrics写入JSON文件或使用可视化工具记录，方便后续分析。

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