NLP——HF-Trainer使用总结

HuggingFace Trainer 整体介绍

  • HuggingFace Trainer 是 transformers 库中提供的一个高级训练 API，旨在简化深度学习模型的训练流程
    • 尤其是对 Transformer 类模型的支持很丰富
  • Trainer 封装了训练循环的核心逻辑，支持多种常见任务（如文本分类、问答、翻译等），并内置了以下多种关键功能：
    • 自动处理单卡/多卡训练、分布式训练
    • 集成日志记录（TensorBoard、W&B等）
    • 支持模型保存、加载和断点续训
    • 内置评估机制，可自定义评估指标
    • 支持早停（Early Stopping）、学习率调度等训练策略
    • 兼容 datasets 库的数据集格式
    • 支持混合精度训练
    • 支持梯度累积（Gradient Accumulation），即在较小的批次上累积梯度，以模拟更大的批次大小
  • HuggingFace Trainer 为 PyTorch 模型提供了完整的训练和评估循环，极大地简化了训练过程，让用户可以专注于模型、数据集和训练参数的配置，而无需手动编写复杂的训练代码
  • Trainer 是一个 开箱即用（out-of-the-box） 的训练工具，它将训练中的各种细节全部封装起来：
    • 前向传播（Forward pass）
    • 计算损失（Loss calculation）
    • 反向传播（Backward pass）
    • 参数更新（Optimizer step）
    • 学习率调整（Learning rate scheduling）
    • 检查点保存（Checkpoint saving）
    • 日志记录（Logging）

对 Trainer 的定制和生态集成

  • Trainer 提供了全面的默认功能，也设计了高度的可定制性
  • 可以通过 子类化（subclassing） 或 重写（overriding） 其内部方法来满足特定的需求，例如：
    • 重写 get_train_dataloader() 方法以自定义数据加载器
    • 重写 compute_loss() 方法以使用自定义的损失函数
    • 重写 compute_metrics() 方法以计算和报告自定义的评估指标
  • Trainer 与 Hugging Face 的其他库（如 datasets 和 accelerate）无缝集成，这使得数据加载、预处理和模型部署变得更加流畅
    • Trainer 特别针对 transformers 库中的模型进行了优化，但也可以与用户自定义的 PyTorch 模型一起使用
  • Trainer 的设计理念是 “开箱即用” ，对于大多数任务，只需要提供 model、args 和 train_dataset 即可
    • 如果需要更精细的控制，可以逐步添加 eval_dataset、data_collator 和 compute_metrics
    • 对于更高级的定制，如自定义训练循环或高级优化策略，则可以利用 callbacks 和 optimizers 等参数

Trainer 类使用示例

安装依赖

  • 首先安装必要的库：
    1
    pip install transformers datasets evaluate accelerate

完整代码示例

  • Trainer 抽象了整个训练循环，只需要提供模型、训练参数、数据集和数据整理器（data_collator），然后调用 trainer.train() 就可以开始训练
  • 以加载一个预训练好的 DistilBERT 模型为例，实现 IMDb 情感分析（正面/负面分类）： 
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    from datasets import load_dataset
    from transformers import (
        AutoTokenizer,
        AutoModelForSequenceClassification,
        TrainingArguments,
        Trainer,
        DataCollatorWithPadding
    )
    import evaluate
    import numpy as np

    dataset = load_dataset("imdb") # 加载数据集（IMDb情感分析：正面/负面分类），包含train和test拆分

    model_name = "distilbert-base-uncased" # 可选择自己的模型
    tokenizer = AutoTokenizer.from_pretrained(model_name) # 加载分词器
    model = AutoModelForSequenceClassification.from_pretrained( # 用 Classification 类加载预训练模型
        model_name,
        num_labels=2  # 二分类任务（正面/负面）
    )
    # Trainer支持多种任务，只需替换对应的模型（如 `AutoModelForQuestionAnswering` 用于问答）和数据集即可

    def preprocess_function(examples): # 数据预处理函数（分词）
        return tokenizer(examples["text"], truncation=True, max_length=512)
    tokenized_dataset = dataset.map(preprocess_function, batched=True) # 对数据集应用预处理

    accuracy = evaluate.load("accuracy") # 定义评估指标（准确率）

    def compute_metrics(eval_pred): # 定义指标计算函数
        predictions, labels = eval_pred
        predictions = np.argmax(predictions, axis=1)  # 取概率最大的类别
        return accuracy.compute(predictions=predictions, references=labels)

    # 重点：配置训练参数（TrainingArguments）
    training_args = TrainingArguments(
        output_dir="./save",                 # 模型保存路径
        learning_rate=2e-5,                  # 学习率
        per_device_train_batch_size=8,       # 每个设备的训练批次大小
        per_device_eval_batch_size=8,        # 每个设备的评估批次大小
        num_train_epochs=3,                  # 训练轮数
        weight_decay=0.01,                   # 权重衰减系数（防止过拟合）
        eval_strategy="epoch",               # 每轮结束后评估，旧版本中该参数名为 eval_strategy，新版本中为 evaluation_strategy
        save_strategy="epoch",               # 每轮结束后保存模型
        load_best_model_at_end=True,         # 训练结束后加载最优模型，确保最终保存的是验证集上表现最好的模型
        logging_dir="./logs",                # 日志保存路径
        logging_steps=100,                   # 每100步记录一次日志
    )

    trainer = Trainer( # 重点：初始化Trainer
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset["train"],  # 训练集
        eval_dataset=tokenized_dataset["test"],    # 评估集
        tokenizer=tokenizer,                       # 传入 Tokenizer
        compute_metrics=compute_metrics,           # 评估指标函数
        data_collator=DataCollatorWithPadding(tokenizer=tokenizer)  # 使用动态padding
    )

    trainer.train() # 重点：开始训练

    # 测试训练结果：使用最优模型进行预测
    sample_text = "This movie is amazing! The acting was incredible and the plot was gripping."
    inputs = tokenizer(sample_text, return_tensors="pt").to(model.device)
    outputs = model(**inputs)
    predicted_class = outputs.logits.argmax().item()
    labels = ["negative", "positive"]
    print(f"预测结果：{labels[predicted_class]}")

启动命令

  • 根据需要启动 trainer_demo.py 的命令需要根据训练需求（如单卡训练、多卡训练、调试模式等）选择不同的方式
基础启动方式（单卡训练）
  • 直接用Python命令运行即可： 
    1
    python trainer_demo.py

多卡训练（使用 accelerate）

  • 为了满足分布式训练，无需修改代码 ，通过 accelerate 工具（已在依赖中安装）启动即可，即 accelerate launch 命令即可启动多卡训练：

    1
    accelerate launch --num_processes=2 trainer_demo.py

  • 更多 accelerate launch 命令启动的用法见：/Notes/PyTorch/PyTorch——HF-Accelerate使用总结

调试模式（单步执行）启动

  • 如需调试代码（如设置断点），可使用 debugpy 或 Python 内置调试器：

    1
    python -m debugpy --wait-for-client --listen 5678 trainer_demo.py
    • 会在5678端口等待调试客户端连接（如 VS Code、PyCharm）
    • 适合排查训练逻辑或数据处理问题

启动混合精读训练

  • 启用 FP16/FP8 混合精度训练，有两种方法

  • 方法一：可在启动命令中指定 ，详情参考 accelerate launch 用法

    1
    accelerate launch --mixed_precision=fp16 trainer_demo.py
    • --mixed_precision 可选值：no（默认）、fp16、bf16、fp8（需特定GPU支持）

  • 方法二：可在 TrainingArguments 中配置：

    1
    2
    3
    4
    training_args = TrainingArguments(
        ...
        fp16=True,  # 启用FP16混合精度
    )
    • 此时用基础命令启动即可：python trainer_demo.py

附录：Trainer 类的详细参数说明

核心参数（通常必须提供）

  • model: 训练或评估的 PyTorch 模型
    • 通常是 transformers 库中继承自 PreTrainedModel 的模型，比如 BertForSequenceClassification
    • 也可以传入一个自定义的 PyTorch 模型，但它需要与 Trainer 的其他组件兼容
  • args: 是一个 TrainingArguments 类的实例
    • 封装了所有的训练超参数，例如学习率、批次大小、训练轮次、日志目录等，是控制训练行为的主要方式
  • train_dataset: 训练数据集
    • 通常是 torch.utils.data.Dataset 或 Hugging Face datasets 库中的 Dataset 对象
    • 提供了这个参数后，就可以调用 trainer.train() 方法开始训练
  • eval_dataset: 评估数据集
    • Trainer 会定期在这个数据集上运行评估，以监控模型性能，并作为最佳模型的选择依据

数据处理相关参数

  • data_collator: 数据整理器
    • 作用是接收数据集中的多个样本，并将它们组合成一个批次（batch）
    • 对于 NLP 任务，它通常会负责对文本序列进行填充（padding）和张量化（tensorizing）
    • 如果不提供，Trainer 会使用一个默认的整理器，但通常只适用于固定长度的输入
    • 对于变长序列，需要提供一个 DataCollatorWithPadding
  • tokenizer: 分词器
    • 用于与 data_collator 协同工作，通常用于处理文本数据
    • 提供分词器可以让 Trainer 在内部自动处理一些数据预处理工作，比如与 DataCollatorWithPadding 结合使用来填充序列
  • model_init: 一个用于初始化模型的函数
    • 如果想使用 Trainer 的超参数搜索（hyperparameter search）功能 ，或者在每次训练前都重新初始化模型，而不是使用固定的 model 实例，就可以使用这个参数
    • 该函数应该不带参数，并返回一个模型实例

训练与评估相关参数

  • compute_metrics: 一个计算评估指标的函数
    • 这个函数接收一个 EvalPrediction 对象（包含预测结果和标签），并返回一个字典，其中包含在评估过程中报告的指标（例如准确率、F1 分数等）
  • callbacks: 一个 TrainerCallback 实例列表
    • 回调函数让你可以在训练循环的特定点（如每个训练步、每个 epoch 结束时）执行自定义逻辑，比如提前停止训练、在 TensorBoard 中记录额外信息等
  • optimizers: 一个包含 优化器和学习率调度器 的元组 (optimizer, lr_scheduler)
    • 这里配置的自定义优化器优先级比 TrainingArguments 中默认配置的更高

其他高级参数

  • preprocess_logits_for_metrics: 一个预处理逻辑的函数
    • 这个函数会在计算指标之前对模型的输出（logits）进行预处理
    • 例如，对于分类任务，你可能希望在计算准确率之前应用 argmax 来获取预测的类别索引
  • args.deepspeed: 这个参数不是直接给 Trainer 的，而是 TrainingArguments 的一部分
    • 可以传入一个 DeepSpeed 配置文件的路径或字典，从而启用 DeepSpeed 进行大规模分布式训练，以实现模型并行、梯度检查点等高级功能

附录：Trainer 中优化器设置

  • 在 HuggingFace Trainer 中设置优化器有两种主要方式：使用内置优化器或自定义优化器

使用内置优化器（简单方式）

  • Trainer 内置了常见优化器（如 Adamw_hf、adamw_torch 等），可通过 TrainingArguments 直接配置，无需额外代码：

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    from transformers import TrainingArguments

    training_args = TrainingArguments(
        output_dir="./saves",
        optim="adamw_torch",  # 优化器设置，可选: adamw_hf, adamw_torch_fused, adafactor等，默认值为 adamw_hf
        learning_rate=2e-5,   # 学习率，默认值 5e-5
        weight_decay=0.01,    # 权重衰减（仅AdamW等支持），默认值为 0 
        adam_beta1=0.9,       # Adam优化器的beta1参数，默认值 0.9
        adam_beta2=0.999,     # Adam优化器的beta2参数，默认值 0.999
        adam_epsilon=1e-8,    # Adam优化器的epsilon参数，默认值 1e-8
    )

  • 常用内置优化器 ：

    • adamw_hf：HuggingFace 实现的 AdamW（默认值）
    • adamw_torch：PyTorch 原生 AdamW
    • adamw_torch_fused：PyTorch 融合版 AdamW（速度更快）
    • adafactor：适合大模型的 Adafactor 优化器（无需设置学习率）

自定义优化器（灵活方式）

  • 如果需要使用 Trainer 不包含的优化器（如RAdam、SGD等），可通过重写 Trainer 类的 create_optimizer 方法实现：

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    class CustomOptimizerTrainer(Trainer):
        def create_optimizer(self): # 重写 Trainer 类
            params = [p for p in self.model.parameters() if p.requires_grad] # 定义需要优化的参数（排除冻结层）
            optimizer = optim.SGD( # 自定义优化器（这里以SGD为例，也可替换为其他优化器）
                params,
                lr=5e-5,       # 学习率
                momentum=0.9,  # SGD动量参数
                weight_decay=0.01
            )
            return optimizer # 返回优化器对象

    training_args = TrainingArguments(
        output_dir="./custom_optimizer_results",
        # ...
        # 注意：自定义优化器时，无需在TrainingArguments中设置optim参数
    )

    # 使用自定义 Trainer 初始化
    trainer = CustomOptimizerTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        tokenizer=tokenizer,
    )

    # 训练
    trainer.train()

  • 常用在需要对不同层设置不同学习率（如微调时固定底层，只训练顶层），可在 create_optimizer 中对参数分组：

    1
    2
    3
    4
    5
    6
    # PyTorch 官方库 `optim` 支持的分不同参数设置不同的学习率
    params = [
        {"params": model.pretrained_layer.parameters(), "lr": 1e-5},  # 底层参数
        {"params": model.classifier.parameters(), "lr": 1e-4}   # 分类头参数
    ]
    optimizer = optim.AdamW(params, weight_decay=0.01)

  • 配合优化器使用时，可通过 TrainingArguments 的 lr_scheduler_type 配置调度策略（如线性衰减、余弦退火等）

  • 自定义优化器需兼容 PyTorch 的优化器接口（继承 torch.optim.Optimizer），否则可能导致训练异常