跳到主要内容

模型训练

本章将详细介绍如何使用 Transformers 的 Trainer API 进行模型训练，包括数据准备、训练配置、评估和保存等完整流程。

Trainer 概述​

什么是 Trainer？​

Trainer 是 Transformers 提供的高级训练 API，封装了训练循环的复杂逻辑，让你只需配置参数即可开始训练。

┌─────────────────────────────────────────────────────────────┐
│                    Trainer 工作流程                          │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   数据准备                                                    │
│      │                                                       │
│      ▼                                                       │
│  ┌─────────────┐                                            │
│  │  Dataset    │  加载和预处理数据                           │
│  │  DataLoader │  批量加载数据                               │
│  └──────┬──────┘                                            │
│         │                                                    │
│   训练循环                                                    │
│         ▼                                                    │
│  ┌─────────────┐    ┌─────────────┐    ┌─────────────┐     │
│  │  Forward    │───▶│  Compute    │───▶│  Backward   │     │
│  │  Pass       │    │  Loss       │    │  & Optimize │     │
│  └─────────────┘    └─────────────┘    └─────────────┘     │
│         │                                                    │
│         ▼                                                    │
│  ┌─────────────┐                                            │
│  │  Evaluation │  定期评估模型性能                           │
│  │  Checkpoint │  保存模型检查点                             │
│  └─────────────┘                                            │
│                                                             │
└─────────────────────────────────────────────────────────────┘

Trainer 的核心功能​

  • 自动混合精度训练：节省显存，加速训练
  • 分布式训练：支持多 GPU 和 DeepSpeed
  • 梯度累积：模拟大批量训练
  • 学习率调度：多种调度策略
  • 日志记录：集成 TensorBoard 和 WandB
  • 模型保存：自动保存检查点

基本训练流程​

1. 准备数据​

from datasets import load_dataset
from transformers import AutoTokenizer

# 加载数据集
dataset = load_dataset("imdb")

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

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

# 应用预处理
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")

print(tokenized_dataset)

2. 加载模型​

from transformers import AutoModelForSequenceClassification

# 加载预训练模型（用于分类任务）
model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=2  # 二分类
)

3. 配置训练参数​

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",           # 输出目录
    evaluation_strategy="epoch",      # 每个 epoch 评估一次
    learning_rate=2e-5,               # 学习率
    per_device_train_batch_size=8,    # 训练批次大小
    per_device_eval_batch_size=8,     # 评估批次大小
    num_train_epochs=3,               # 训练轮数
    weight_decay=0.01,                # 权重衰减
    save_strategy="epoch",            # 每个 epoch 保存一次
    load_best_model_at_end=True,      # 训练结束时加载最佳模型
    logging_dir="./logs",             # 日志目录
    logging_steps=10,                 # 每 10 步记录一次日志
    report_to="tensorboard",          # 报告到 TensorBoard
)

4. 创建 Trainer 并开始训练​

from transformers import Trainer, DataCollatorWithPadding

# 数据整理器（用于动态填充）
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

# 创建 Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"].shuffle(seed=42).select(range(1000)),
    eval_dataset=tokenized_dataset["test"].shuffle(seed=42).select(range(500)),
    tokenizer=tokenizer,
    data_collator=data_collator,
)

# 开始训练
trainer.train()

5. 保存和评估​

# 保存模型
trainer.save_model("./my_trained_model")

# 评估模型
eval_results = trainer.evaluate()
print(f"评估结果: {eval_results}")

# 推送到 Hugging Face Hub
trainer.push_to_hub("my-username/my-model")

TrainingArguments 详解​

基本参数​

from transformers import TrainingArguments

training_args = TrainingArguments(
    # 输出和日志
    output_dir="./results",           # 模型和检查点保存目录
    logging_dir="./logs",             # TensorBoard 日志目录
    logging_steps=100,                # 每 N 步记录日志
    report_to="tensorboard",          # 日志报告目标
    
    # 训练参数
    num_train_epochs=3,               # 训练轮数
    learning_rate=5e-5,               # 初始学习率
    weight_decay=0.01,                # 权重衰减系数
    warmup_ratio=0.1,                 # 预热步数比例
    
    # 批次参数
    per_device_train_batch_size=8,    # 每个设备的训练批次
    per_device_eval_batch_size=8,     # 每个设备的评估批次
    gradient_accumulation_steps=4,    # 梯度累积步数
    
    # 评估和保存
    evaluation_strategy="steps",      # 评估策略: "no", "steps", "epoch"
    eval_steps=500,                   # 每 N 步评估
    save_strategy="steps",            # 保存策略
    save_steps=500,                   # 每 N 步保存
    save_total_limit=3,               # 最多保存的检查点数量
    load_best_model_at_end=True,      # 结束时加载最佳模型
    metric_for_best_model="eval_loss", # 选择最佳模型的指标
    
    # 优化器
    optim="adamw_torch",              # 优化器类型
    lr_scheduler_type="linear",       # 学习率调度器
)

高级参数​

training_args = TrainingArguments(
    # 混合精度训练
    fp16=True,                        # 启用 FP16 混合精度
    bf16=True,                        # 启用 BF16 混合精度（Ampere GPU）
    
    # 梯度相关
    max_grad_norm=1.0,                # 梯度裁剪阈值
    gradient_checkpointing=True,      # 启用梯度检查点
    
    # 分布式训练
    local_rank=-1,                    # 本地进程排名
    ddp_find_unused_parameters=False, # DDP 参数
    deepspeed="ds_config.json",       # DeepSpeed 配置文件
    
    # 其他
    seed=42,                          # 随机种子
    dataloader_num_workers=4,         # 数据加载器工作进程数
    remove_unused_columns=True,       # 移除未使用的列
)

数据准备详解​

使用 Hugging Face Datasets​

from datasets import load_dataset

# 加载内置数据集
dataset = load_dataset("glue", "sst2")

# 加载本地数据文件
# dataset = load_dataset("csv", data_files={"train": "train.csv", "test": "test.csv"})

# 查看数据集结构
print(dataset)
print(dataset["train"][0])

数据预处理​

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def tokenize_function(examples):
    """分词处理函数"""
    return tokenizer(
        examples["sentence"],
        padding="max_length",
        truncation=True,
        max_length=128
    )

# 应用分词
tokenized_datasets = dataset.map(
    tokenize_function,
    batched=True,
    num_proc=4,  # 使用 4 个进程并行处理
    remove_columns=["sentence", "idx"]  # 移除原始列
)

# 重命名标签列
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")

# 设置格式为 PyTorch 张量
tokenized_datasets.set_format("torch")

数据整理器（Data Collator）​

from transformers import DataCollatorWithPadding, DataCollatorForLanguageModeling

# 动态填充（推荐）
data_collator = DataCollatorWithPadding(
    tokenizer=tokenizer,
    padding=True,  # 动态填充到批次最大长度
    max_length=None,
)

# 用于语言建模的数据整理器（MLM）
mlm_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=True,           # 启用掩码语言建模
    mlm_probability=0.15,  # 掩码比例
)

自定义训练逻辑​

自定义损失函数​

from transformers import Trainer
import torch.nn as nn

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        """自定义损失计算"""
        labels = inputs.get("labels")
        
        # 前向传播
        outputs = model(**inputs)
        logits = outputs.get("logits")
        
        # 自定义损失（例如：加权交叉熵）
        loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 2.0]).to(logits.device))
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        
        return (loss, outputs) if return_outputs else loss

# 使用自定义 Trainer
trainer = CustomTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

自定义评估指标​

import numpy as np
from sklearn.metrics import accuracy_score, precision_recall_fscore_support

def compute_metrics(eval_pred):
    """自定义评估指标"""
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    
    # 计算指标
    accuracy = accuracy_score(labels, predictions)
    precision, recall, f1, _ = precision_recall_fscore_support(
        labels, predictions, average="binary"
    )
    
    return {
        "accuracy": accuracy,
        "precision": precision,
        "recall": recall,
        "f1": f1,
    }

# 在 Trainer 中使用
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,
)

自定义回调函数​

from transformers import TrainerCallback

class EarlyStoppingCallback(TrainerCallback):
    """早停回调"""
    def __init__(self, patience=3):
        self.patience = patience
        self.best_loss = float('inf')
        self.counter = 0
    
    def on_evaluate(self, args, state, control, metrics=None, **kwargs):
        current_loss = metrics.get("eval_loss", float('inf'))
        
        if current_loss < self.best_loss:
            self.best_loss = current_loss
            self.counter = 0
        else:
            self.counter += 1
            if self.counter >= self.patience:
                control.should_training_stop = True
                print(f"Early stopping triggered after {state.global_step} steps")
        
        return control

# 使用回调
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[EarlyStoppingCallback(patience=3)],
)

分布式训练​

多 GPU 训练​

# 使用 torchrun 启动分布式训练
torchrun --nproc_per_node=4 train.py

# 或在代码中设置
import torch.distributed as dist
# TrainingArguments 配置
training_args = TrainingArguments(
    # ... 其他参数
    local_rank=-1,  # 自动检测
    ddp_find_unused_parameters=False,
)

DeepSpeed 集成​

# ds_config.json
{
    "fp16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
    },
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
}
training_args = TrainingArguments(
    # ... 其他参数
    deepspeed="ds_config.json",
)

训练技巧​

学习率调度​

training_args = TrainingArguments(
    learning_rate=5e-5,
    lr_scheduler_type="cosine",  # "linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"
    warmup_ratio=0.1,  # 预热步数比例
    num_train_epochs=3,
)

梯度累积​

training_args = TrainingArguments(
    per_device_train_batch_size=2,    # 小批次
    gradient_accumulation_steps=8,    # 累积 8 步
    # 等效批次大小 = 2 * 8 = 16
)

混合精度训练​

training_args = TrainingArguments(
    fp16=True,    # 启用 FP16（V100, RTX 系列）
    # 或
    bf16=True,    # 启用 BF16（A100, H100）
)

完整示例​

#!/usr/bin/env python3
"""
完整的文本分类训练示例
"""

from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
)
import numpy as np
from sklearn.metrics import accuracy_score

def main():
    # 1. 加载数据
    print("Loading dataset...")
    dataset = load_dataset("imdb")
    
    # 2. 加载分词器和模型
    print("Loading tokenizer and model...")
    model_name = "distilbert-base-uncased"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
    
    # 3. 预处理数据
    print("Preprocessing data...")
    def preprocess(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            padding=False,
            max_length=512
        )
    
    tokenized = dataset.map(preprocess, batched=True, remove_columns=["text"])
    tokenized = tokenized.rename_column("label", "labels")
    tokenized.set_format("torch")
    
    # 4. 配置训练参数
    training_args = TrainingArguments(
        output_dir="./imdb_classifier",
        evaluation_strategy="epoch",
        learning_rate=2e-5,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=16,
        num_train_epochs=3,
        weight_decay=0.01,
        save_strategy="epoch",
        load_best_model_at_end=True,
        logging_steps=100,
        fp16=True,
    )
    
    # 5. 定义评估指标
    def compute_metrics(eval_pred):
        predictions, labels = eval_pred
        predictions = np.argmax(predictions, axis=1)
        return {"accuracy": accuracy_score(labels, predictions)}
    
    # 6. 创建 Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized["train"].shuffle(seed=42).select(range(5000)),
        eval_dataset=tokenized["test"].shuffle(seed=42).select(range(1000)),
        tokenizer=tokenizer,
        data_collator=DataCollatorWithPadding(tokenizer),
        compute_metrics=compute_metrics,
    )
    
    # 7. 训练
    print("Starting training...")
    trainer.train()
    
    # 8. 评估
    print("Evaluating...")
    results = trainer.evaluate()
    print(f"Final results: {results}")
    
    # 9. 保存
    trainer.save_model("./imdb_classifier_final")
    print("Model saved!")

if __name__ == "__main__":
    main()

下一步​

掌握 Trainer 的使用后，你可以继续学习：

参考资源​