Skip to content

训练流程Trainer

手写训练循环(forward → loss → backward → step)不难,但一旦涉及 GPU、混合精度、梯度累积、断点续训、日志……就会很繁琐。Trainer 把这些工程细节一站式封装好,你只需描述「训什么、怎么训」。

1. 训练一个模型需要哪些要素

┌─────────────────────────────────────────────────┐
│  Trainer 把下面这些拼在一起:                        │
│   model          要训练的模型                      │
│   args           超参(学习率/batch/轮数/精度...)    │
│   train_dataset  训练数据(上一页处理好的)           │
│   eval_dataset   验证数据                          │
│   data_collator  如何把样本拼成 batch              │
│   compute_metrics 评估指标(可选)                   │
└─────────────────────────────────────────────────┘

2. 一个完整的微调示例(情感分类)

python
import numpy as np
from datasets import load_dataset
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    DataCollatorWithPadding, TrainingArguments, Trainer,
)

# ① 数据
ds = load_dataset("imdb").train_test_split(test_size=0.05, seed=42)
tok = AutoTokenizer.from_pretrained("distilbert-base-uncased")
ds = ds.map(lambda b: tok(b["text"], truncation=True, max_length=256), batched=True)

# ② 模型
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2)

# ③ 训练超参
args = TrainingArguments(
    output_dir="./out",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    num_train_epochs=2,
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    fp16=True,                 # 混合精度, 省显存提速(GPU)
    warmup_ratio=0.1,          # 学习率预热
    logging_steps=50,
)

# ④ 指标
def acc(eval_pred):
    logits, labels = eval_pred
    return {"accuracy": (logits.argmax(-1) == labels).mean()}

# ⑤ 组装并开训
trainer = Trainer(
    model=model, args=args,
    train_dataset=ds["train"], eval_dataset=ds["test"],
    data_collator=DataCollatorWithPadding(tokenizer=tok),
    compute_metrics=acc,
)
trainer.train()                # 一行启动整个训练循环
print(trainer.evaluate())      # 验证集指标

跑完后,trainer.save_model("./out/final") 即可保存,之后像普通模型一样加载使用。

3. TrainingArguments 高频参数

参数作用
learning_rate学习率(微调 BERT 类常用 1e-5 ~ 5e-5)
per_device_train_batch_size每卡 batch;显存不够就调小
gradient_accumulation_steps梯度累积,用时间换「等效大 batch」
num_train_epochs / max_steps训练时长
fp16 / bf16混合精度(GPU 强烈建议开)
warmup_ratio学习率线性预热比例,稳定早期训练
eval_strategy / save_strategy何时评估/存档
load_best_model_at_end训练结束回载验证最优的 checkpoint
weight_decay权重衰减(正则)
gradient_checkpointing用重算换显存,训练更大模型

4. Trainer 帮你自动做了什么

✓ 设备放置(把模型/数据搬到 GPU)、DDP 多卡
✓ 混合精度 autocast + GradScaler
✓ 梯度累积 / 梯度裁剪 / 学习率调度(带 warmup)
✓ 周期性 eval / save / logging
✓ 断点续训: trainer.train_from_checkpoint(...)

5. 何时不用 Trainer:手写循环

学习原理、或需要完全自定义(如 NanoGPT)时,手写循环更透明。核心就五步:

python
opt = torch.optim.AdamW(model.parameters(), lr=6e-4)
for x, y in loader:
    x, y = x.to(device), y.to(device)
    logits = model(x)                    # ① 前向
    loss = nn.functional.cross_entropy(logits.view(-1, V), y.view(-1))  # ② 损失
    opt.zero_grad()                      # ③ 清梯度
    loss.backward()                      # ④ 反向传播
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 梯度裁剪
    opt.step()                           # ⑤ 更新参数

这段正是 训练与调参 里我们会逐行讲透、并跑通的循环。理解手写循环,才能明白 Trainer 帮你做了什么。

小结

  • Trainer = 模型 + TrainingArguments + 数据集 + collator 的一站式训练器。
  • 微调 BERT 类模型,学习率小、配 warmup、开混合精度是常规操作。
  • gradient_accumulation/gradient_checkpointing 是显存不足时的救命稻草。
  • 想真正吃透,务必会写 5 步式手写训练循环——实战章节见。

下一步

工具链齐了,进入全书高潮:从零训练一个能说话的 Transformer 小模型实战概览