transformers 的Trainer的用法
Trainer 自动处理以下任务:训练循环:自动实现 epoch 迭代、批次加载优化器&学习率调度:内置 AdamW 并支持自定义分布式训练:自动支持单机多卡(DataParallel/DistributedDataParallel)混合精度训练:通过 fp16=True 启用日志记录:集成 TensorBoard、Weights & Biases 等模型保
·
transformers 库中的 Trainer
类是一个高级 API,它简化了训练和评估 transformer 模型的流程。下面我将从核心概念、基本用法到高级技巧进行全面讲解:
1. 核心功能
Trainer 自动处理以下任务:
- 训练循环:自动实现 epoch 迭代、批次加载
- 优化器&学习率调度:内置 AdamW 并支持自定义
- 分布式训练:自动支持单机多卡(DataParallel/DistributedDataParallel)
- 混合精度训练:通过 fp16=True 启用
- 日志记录:集成 TensorBoard、Weights & Biases 等
- 模型保存:定期保存检查点 + 最佳模型保存
- 评估指标计算:自动计算验证集指标
2. 基础使用步骤
步骤 1:准备组件
from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer
from datasets import load_dataset
# 加载模型
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
# 加载数据集
dataset = load_dataset("glue", "mrpc")
train_set = dataset["train"]
eval_set = dataset["validation"]
# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
# 数据预处理函数
def tokenize(batch):
return tokenizer(batch["sentence1"], batch["sentence2"], padding="max_length", truncation=True)
train_set = train_set.map(tokenize, batched=True)
eval_set = eval_set.map(tokenize, batched=True)
步骤 2:配置训练参数
training_args = TrainingArguments(
output_dir="./results", # 输出目录
num_train_epochs=3, # 训练轮数
per_device_train_batch_size=16, # 每个GPU的批次大小
per_device_eval_batch_size=64, # 验证批次大小
evaluation_strategy="epoch", # 每个epoch后评估
save_strategy="epoch", # 每个epoch后保存模型
learning_rate=2e-5, # 学习率
weight_decay=0.01, # 权重衰减
fp16=True, # 混合精度训练
logging_dir="./logs", # 日志目录
report_to="tensorboard", # 日志工具
load_best_model_at_end=True, # 训练结束加载最佳模型
metric_for_best_model="accuracy", # 最佳模型指标
)
步骤 3:自定义评估指标
import numpy as np
from sklearn.metrics import accuracy_score
def compute_metrics(eval_pred):
logits, labels = eval_pred
predictions = np.argmax(logits, axis=-1)
acc = accuracy_score(labels, predictions)
return {"accuracy": acc}
步骤 4:实例化 Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_set,
eval_dataset=eval_set,
compute_metrics=compute_metrics, # 自定义评估函数
)
步骤 5:开始训练
trainer.train()
步骤 6:评估与预测
results = trainer.evaluate()
print(results)
# 对新数据预测
test_inputs = tokenizer("Hello world!", return_tensors="pt")
predictions = trainer.predict(test_inputs)
3. 高级特性与技巧
3.1 自定义优化器
from torch.optim import AdamW
from transformers import get_scheduler
optimizer = AdamW(model.parameters(), lr=2e-5)
lr_scheduler = get_scheduler(
"linear",
optimizer=optimizer,
num_warmup_steps=100,
num_training_steps=1000
)
trainer = Trainer(
optimizers=(optimizer, lr_scheduler),
... # 其他参数
)
3.2 回调函数(Callbacks)
from transformers import TrainerCallback
class CustomCallback(TrainerCallback):
def on_epoch_end(self, args, state, control, **kwargs):
print(f"Epoch {state.epoch} finished!")
trainer.add_callback(CustomCallback())
3.3 恢复训练
trainer.train(resume_from_checkpoint=True) # 自动加载最新检查点
# 或指定具体路径
# trainer.train(resume_from_checkpoint="/path/to/checkpoint")
3.4 梯度累积与裁剪
在 TrainingArguments 中设置:
gradient_accumulation_steps=4, # 每4个步骤更新一次权重
max_grad_norm=1.0 # 梯度裁剪阈值
3.5 类权重(不平衡数据)
from torch import nn
class WeightedTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False):
labels = inputs.pop("labels")
outputs = model(**inputs)
loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 2.0])) # 类权重
loss = loss_fct(outputs.logits, labels)
return (loss, outputs) if return_outputs else loss
4. 常见问题解决
4.1 内存不足(OOM)
- 降低 per_device_train_batch_size
- 启用梯度累积 (gradient_accumulation_steps)
- 使用 fp16 混合精度
- 尝试梯度检查点(在模型配置中设置 gradient_checkpointing=True)
4.2 自定义数据加载
继承 Trainer 并重写 get_train_dataloader 方法:
class CustomTrainer(Trainer):
def get_train_dataloader(self):
return DataLoader(..., collate_fn=custom_collate)
4.3 多标签分类
需要自定义损失函数:
class MultiLabelTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False):
labels = inputs.pop("labels")
outputs = model(**inputs)
logits = outputs.logits
loss_fct = nn.BCEWithLogitsLoss()
loss = loss_fct(logits, labels.float())
return (loss, outputs) if return_outputs else loss
5. 完整工作流示例
# 安装依赖:pip install transformers datasets torch
from transformers import AutoModel, AutoTokenizer, TrainingArguments, Trainer
from datasets import load_dataset
import torch
# 1. 加载数据与模型
dataset = load_dataset("imdb")
model_name = "bert-base-uncased"
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# 2. 数据预处理
def tokenize(examples):
return tokenizer(examples["text"], padding="max_length", truncation=True)
dataset = dataset.map(tokenize, batched=True)
# 3. 设置训练参数
training_args = TrainingArguments(
output_dir="./results",
num_train_epochs=3,
per_device_train_batch_size=8,
evaluation_strategy="epoch"
)
# 4. 实例化Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset["train"],
eval_dataset=dataset["test"],
compute_metrics=lambda p: {"accuracy": (p.predictions.argmax(-1) == p.label_ids).mean()}
)
# 5. 训练与评估
trainer.train()
results = trainer.evaluate()
print(f"Final results: {results}")
6. 关键注意事项
- 数据格式:确保数据集包含 input_ids, attention_mask 和 labels 字段
- 标签处理:分类任务标签必须是 0 开始的连续整数
- 设备兼容:Trainer 自动使用 GPU(若可用)
- 模型保存:最佳模型保存在 output_dir 下的 best_model 目录
- 日志查看:使用 tensorboard --logdir=./logs 查看训练曲线
通过Trainer
,80% 的常见训练需求可通过配置解决,剩余 20% 可通过继承重写实现。官方文档提供了更多进阶用例:HuggingFace Trainer 文档
更多推荐
所有评论(0)