Huggingface 核心模块(四): Trainer
huggingface
本文字数:951 字 | 阅读时长 ≈ 4 min

Huggingface 核心模块(四): Trainer

huggingface
本文字数:951 字 | 阅读时长 ≈ 4 min

1. Trainer 快速入门

1.1 基础示例

让我们通过一个简单的文本分类示例来了解 Trainer 的基本用法:

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
import numpy as np
import evaluate

# 加载数据集和评估指标
dataset = load_dataset("yelp_review_full")
metric = evaluate.load("accuracy")

# 初始化tokenizer和预处理函数
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased")

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

# 数据预处理
tokenized_datasets = dataset.map(tokenize_function, batched=True)
small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))

# 模型训练设置
model = AutoModelForSequenceClassification.from_pretrained("google-bert/bert-base-cased", num_labels=5)
training_args = TrainingArguments(output_dir="test_trainer", evaluation_strategy="epoch")

# 初始化 Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train_dataset,
    eval_dataset=small_eval_dataset,
    compute_metrics=compute_metrics,
)

# 开始训练
trainer.train()

1.2 数据集访问说明

在使用 Huggingface datasets 时,可以通过以下方式访问具体样本:

# 访问训练集中的第101个样本
sample = dataset["train"][100]
print(sample)

# 输出示例:
{
    'label': 0, 
    'text': 'My expectations for McDonalds are t rarel...'
}

2. Trainer 核心组件解析

2.1 训练流程详解

trainer.train() 方法执行了完整的模型训练流程,主要包含以下阶段:

初始化阶段

训练循环

训练结束处理

2.2 数据加载器详解

Trainer 的数据加载过程主要通过 get_train_dataloader 方法实现:

def get_train_dataloader(self) -> DataLoader:
    """返回训练用的 DataLoader"""
    if self.train_dataset is None:
        raise ValueError("Trainer: training requires a train_dataset.")

    train_dataset = self.train_dataset
    data_collator = self.data_collator

    # 数据集列处理
    if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
        train_dataset = self._remove_unused_columns(train_dataset, description="training")
    
    # DataLoader 参数配置
    dataloader_params = {
        "batch_size": self._train_batch_size,
        "collate_fn": data_collator,
        "num_workers": self.args.dataloader_num_workers,
        "pin_memory": self.args.dataloader_pin_memory,
    }

    # 返回准备好的 DataLoader
    return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))

3. 高级特性:数据采样策略

3.1 随机采样器 vs 长度分组采样器

Trainer 支持两种主要的采样策略:

RandomSampler

# 完全随机打乱数据
sentences = [
    "我喜欢深度学习",           # 长度: 6
    "这是一个例子",           # 长度: 5
    "自然语言处理非常有趣",     # 长度: 9
    "transformer很强大",      # 长度: 11
]

# 可能的批次分组(需要较多padding):
Batch 1: ["transformer很强大", "我喜欢深度学习"]        # 长度: 11, 6
Batch 2: ["自然语言处理非常有趣", "这是一个例子"]        # 长度: 9, 5

LengthGroupedSampler

# 相似长度的句子会被分到一起
# 可能的批次分组(最小化padding):
Batch 1: ["这是一个例子", "我喜欢深度学习"]           # 长度: 5, 6
Batch 2: ["自然语言处理非常有趣", "transformer很强大"]  # 长度: 9, 11

3.2 分布式训练中的采样器

在分布式训练环境中,Trainer 会自动将 RandomSampler 转换为 DistributedSampler:

# 单GPU训练
trainer = Trainer(
    model=model,
    args=TrainingArguments("output_dir"),
    train_dataset=dataset
)
# → 使用 RandomSampler

# 多GPU训练
trainer = Trainer(
    model=model,
    args=TrainingArguments(
        "output_dir",
        per_device_train_batch_size=8,
        local_rank=local_rank,
    ),
    train_dataset=dataset
)
# → RandomSampler 自动转换为 DistributedSampler

4. 最佳实践建议

  1. 对于变长序列任务,建议启用长度分组采样:
training_args = TrainingArguments(
    "output_dir",
    per_device_train_batch_size=16,
    group_by_length=True  # 启用长度分组
)
  1. 分布式训练无需手动配置采样器,Trainer 会自动处理

  2. 根据任务特点选择合适的评估策略:

training_args = TrainingArguments(
    "output_dir",
    evaluation_strategy="steps",  # 或 "epoch"
    eval_steps=500,              # 每500步评估一次
)
  1. 合理设置批次大小和学习率:
training_args = TrainingArguments(
    "output_dir",
    per_device_train_batch_size=8,
    learning_rate=2e-5,
    gradient_accumulation_steps=4,  # 累积梯度
)