Appearance
训练流程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 小模型 → 实战概览。