如何在DistilBERT文本分类训练过程中实现交叉验证

实现思路 #

• 首先拆分出独立的最终测试集，全程不参与交叉验证过程，用来评估最终泛化性能
• 采用K折交叉验证在剩余数据集上完成训练调参，每折都会重新初始化模型、拆分训练/验证集，避免数据泄露
• 保留每折的训练/验证指标，最后取平均得到更稳定的性能评估结果

完整可运行代码 #

import tensorflow as tf
from sklearn.model_selection import train_test_split, KFold
from transformers import TFDistilBertForSequenceClassification, DistilBertTokenizer
import numpy as np

# 超参数设置，可根据自己的需求调整
K = 5 # 交叉验证折数
EPOCHS = 10
BATCH_SIZE = 16
RANDOM_STATE = 42
LEARNING_RATE = 2e-5 # BERT类模型微调推荐学习率区间为2e-5~5e-5

# 第一步：先拆分出独立的最终测试集，全程不参与训练/验证过程
X_train_val, X_test, y_train_val, y_test = train_test_split(X, y, test_size=0.1, random_state=RANDOM_STATE)

# 初始化分词器
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
# 提前编码测试集，后续直接使用
test_encoded = tokenizer(X_test, truncation=True, padding=True, return_tensors="tf")
test_data = tf.data.Dataset.from_tensor_slices((dict(test_encoded), y_test)).batch(BATCH_SIZE)

# 初始化K折交叉验证器
kf = KFold(n_splits=K, shuffle=True, random_state=RANDOM_STATE)
# 存储每折的评估指标
fold_val_acc = []
fold_val_loss = []

for fold, (train_idx, val_idx) in enumerate(kf.split(X_train_val)):
    print(f"########## 开始训练第{fold+1}/{K}折 ##########")
    # 拆分当前折的训练和验证数据
    X_train_fold = X_train_val[train_idx]
    y_train_fold = y_train_val[train_idx]
    X_val_fold = X_train_val[val_idx]
    y_val_fold = y_train_val[val_idx]

    # 编码当前折的训练数据
    train_encoded = tokenizer(X_train_fold.tolist(), truncation=True, padding=True, return_tensors="tf")
    train_data = tf.data.Dataset.from_tensor_slices((dict(train_encoded), y_train_fold)).shuffle(1000).batch(BATCH_SIZE)

    # 编码当前折的验证数据
    val_encoded = tokenizer(X_val_fold.tolist(), truncation=True, padding=True, return_tensors="tf")
    val_data = tf.data.Dataset.from_tensor_slices((dict(val_encoded), y_val_fold)).batch(BATCH_SIZE)

    # 每折重新初始化模型，避免权重泄露
    model = TFDistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=2)
    optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE)
    model.compile(optimizer=optimizer, loss=model.compute_loss, metrics=['accuracy'])

    # 训练当前折模型
    history = model.fit(
        train_data,
        validation_data=val_data,
        epochs=EPOCHS,
        verbose=1
    )

    # 保存当前折的最优验证指标
    fold_val_acc.append(max(history.history['val_accuracy']))
    fold_val_loss.append(min(history.history['val_loss']))
    print(f"第{fold+1}折验证准确率：{fold_val_acc[-1]:.4f}，验证损失：{fold_val_loss[-1]:.4f}\n")

# 输出交叉验证平均结果
print(f"########## {K}折交叉验证平均结果 ##########")
print(f"平均验证准确率：{np.mean(fold_val_acc):.4f} ± {np.std(fold_val_acc):.4f}")
print(f"平均验证损失：{np.mean(fold_val_loss):.4f} ± {np.std(fold_val_loss):.4f}")

# 可选：训练最终模型用全部train_val数据训练，再在测试集评估泛化性能
final_model = TFDistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=2)
final_model.compile(optimizer=optimizer, loss=final_model.compute_loss, metrics=['accuracy'])
train_val_encoded = tokenizer(X_train_val.tolist(), truncation=True, padding=True, return_tensors="tf")
train_val_data = tf.data.Dataset.from_tensor_slices((dict(train_val_encoded), y_train_val)).shuffle(1000).batch(BATCH_SIZE)
final_model.fit(train_val_data, epochs=EPOCHS, verbose=1)
test_loss, test_acc = final_model.evaluate(test_data, verbose=0)
print(f"最终测试集准确率：{test_acc:.4f}，测试损失：{test_loss:.4f}")

注意事项 #

• 原代码中test_data和test_encoded的定义顺序写反了，上述代码已经做了修正
• 每折必须重新初始化模型，不能复用之前折训练后的权重，否则会出现严重的数据泄露，导致评估结果虚高
• 如果数据集较小可以将K设为10，数据集较大可以设为3或5减少训练耗时

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