如何在Huggingface Trainer中实现K-Fold交叉验证？

Huggingface Trainer没有内置的K-Fold交叉验证功能，需要手动结合数据拆分工具（如scikit-learn的KFold）来实现，核心思路是循环划分K份数据集，每份分别作为验证集，其余作为训练集，在每个fold中独立初始化模型并完成微调与评估，最后汇总所有fold的指标得到更稳健的结果。

具体实现步骤与代码示例 #

1. 导入依赖并准备数据 #

from datasets import load_dataset
from sklearn.model_selection import KFold
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    EvalPrediction
)
import numpy as np

# 加载文本分类数据集（示例用IMDB数据集，替换为你的数据集即可）
dataset = load_dataset("imdb")["train"]
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 数据预处理函数
def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)

tokenized_dataset = dataset.map(preprocess_function, batched=True)
tokenized_dataset = tokenized_dataset.remove_columns(["text"])
tokenized_dataset = tokenized_dataset.rename_column("label", "labels")
tokenized_dataset.set_format("torch")

2. 初始化K-Fold拆分器并循环训练 #

# 配置5折交叉验证参数，shuffle=True保证数据打乱
kf = KFold(n_splits=5, shuffle=True, random_state=42)
all_eval_metrics = []

# 遍历每个fold
for fold_num, (train_idx, val_idx) in enumerate(kf.split(tokenized_dataset)):
    print(f"===== 开始训练第 {fold_num+1} 折 =====")
    
    # 拆分当前fold的训练集与验证集
    train_dataset = tokenized_dataset.select(train_idx)
    val_dataset = tokenized_dataset.select(val_idx)
    
    # 重新初始化模型：必须每个fold都加载全新预训练模型，避免参数泄露
    model = AutoModelForSequenceClassification.from_pretrained(
        "bert-base-uncased",
        num_labels=2  # 根据你的分类任务调整类别数量
    )
    
    # 配置训练参数
    training_args = TrainingArguments(
        output_dir=f"./results_fold_{fold_num+1}",  # 每个fold单独的输出目录
        per_device_train_batch_size=16,
        per_device_eval_batch_size=16,
        num_train_epochs=3,
        logging_dir=f"./logs_fold_{fold_num+1}",
        logging_steps=100,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True
    )
    
    # 定义评估指标（以准确率为例）
    def compute_metrics(p: EvalPrediction):
        preds = np.argmax(p.predictions, axis=1)
        return {"accuracy": (preds == p.label_ids).mean()}
    
    # 初始化Trainer并开始训练
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        compute_metrics=compute_metrics
    )
    
    trainer.train()
    fold_metrics = trainer.evaluate()
    all_eval_metrics.append(fold_metrics["eval_accuracy"])
    print(f"第 {fold_num+1} 折验证准确率: {fold_metrics['eval_accuracy']:.4f}")

# 计算所有fold的平均指标
avg_accuracy = np.mean(all_eval_metrics)
print(f"\n5折交叉验证平均准确率: {avg_accuracy:.4f}")

关键注意事项 #

• 每次fold重新初始化模型：必须在每个循环中重新加载预训练模型，避免前一个fold的训练参数污染当前fold，导致数据泄露。
• 独立的输出与日志目录：为每个fold设置单独的output_dir和logging_dir，防止不同fold的训练结果、日志互相覆盖。
• 固定随机种子：设置random_state保证K-Fold拆分的可重复性，训练时也建议固定全局随机种子，提升结果稳定性。

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