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 文档
Logo

技术共进,成长同行——讯飞AI开发者社区

更多推荐