使用 Hugging Face Trainer 微调小型模型

1. 概述

本文将介绍如何使用 transformers 库的 Trainer API 来微调一个小型模型（DistilBERT），用于产品评论情感分类。学完本文后，你将了解如何准备数据集、tokenize、配置 TrainingArguments、训练、评估和保存模型以供推理使用。

贯穿本文的例子：OrderReview—客户（Customer）对订单（Order）的评价，标注为 positive/negative。

2. 环境准备

安装必要的库。你不需要强大的 GPU，小型模型如 DistilBERT 可以在 CPU 上处理小数据集，虽然速度会慢一些。

pip install transformers datasets evaluate accelerate torch scikit-learn

注意： 如果你的机器有 GPU，安装对应的 CUDA 版本 PyTorch 会大幅加速训练。使用 CPU 时，应该限制数据行数和 epoch 数量，以免演示花费过长时间。

3. 准备数据集

与其加载大型数据集，我们创建一个小的数据集来模拟订单评论，便于理解。在实际应用中，你可以用 load_dataset("imdb") 或自己的 CSV 数据替换。

from datasets import Dataset

reviews = [
    {"text": "订单配送迅速，包装仔细，非常满意", "label": 1},
    {"text": "产品有缺陷，配送缓慢，不符合描述", "label": 0},
    {"text": "质量好，价格合理，会再次购买", "label": 1},
    {"text": "服务差，客户支持缓慢", "label": 0},
    {"text": "非常满意这个订单，五星好评", "label": 1},
    {"text": "收到时货物受损，很失望", "label": 0},
]

dataset = Dataset.from_list(reviews)
dataset = dataset.train_test_split(test_size=0.33, seed=42)
print(dataset)

DatasetDict({
    train: Dataset({features: ['text', 'label'], num_rows: 4})
    test: Dataset({features: ['text', 'label'], num_rows: 2})
})

Dataset.from_list 将列表字典转换为 Hugging Face 的 Dataset。train_test_split 将数据分为训练/测试集—对于真实数据集，你应该有数千行或更多，这里仅作为演示工作流。

4. 数据 Tokenize

模型无法直接理解文本，需要 tokenizer 将其转换为 input_ids 和 attention_mask。

from transformers import AutoTokenizer

checkpoint = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

def tokenize_fn(batch):
    return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=32)

tokenized_dataset = dataset.map(tokenize_fn, batched=True)
print(tokenized_dataset["train"][0].keys())

dict_keys(['text', 'label', 'input_ids', 'attention_mask'])

• truncation=True: 截断超过 max_length 的句子。
• padding="max_length": 填充 token 使所有句子长度相同，便于批量处理。
• map(..., batched=True): 按批次应用 tokenize 函数，比逐行处理快。

注意： 英文 DistilBERT 在实际应用中不适合微调越南文评论。这里用它来演示技术；对于越南文数据，应该选择像 vinai/phobert-base 这样的 checkpoint。

5. 加载模型和定义标签

from transformers import AutoModelForSequenceClassification

label2id = {"negative": 0, "positive": 1}
id2label = {v: k for k, v in label2id.items()}

model = AutoModelForSequenceClassification.from_pretrained(
    checkpoint,
    num_labels=2,
    id2label=id2label,
    label2id=label2id,
)

num_labels=2 在预训练的编码器上创建一个新的分类头，有 2 个输出（positive/negative）。id2label/label2id 帮助模型输出返回可读的标签而不是数字 0/1。

6. 配置 TrainingArguments

TrainingArguments 汇总了所有超参数和训练循环配置。

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./order-review-model",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=3,
    weight_decay=0.01,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    logging_steps=1,
)

• eval_strategy="epoch": 每个 epoch 后评估一次，而不是每个 step。
• load_best_model_at_end=True: 训练完成后，Trainer 自动加载根据 metric_for_best_model 选出的最佳 checkpoint。
• per_device_train_batch_size=4: 批次较小，因为演示数据集也小；对于真实数据，可根据 VRAM 增加到 16/32。

7. 定义度量函数

Trainer 需要一个 compute_metrics 函数来定义如何评估模型。

import numpy as np
import evaluate

accuracy = evaluate.load("accuracy")
f1 = evaluate.load("f1")

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    acc = accuracy.compute(predictions=predictions, references=labels)
    f1_score = f1.compute(predictions=predictions, references=labels)
    return {"accuracy": acc["accuracy"], "f1": f1_score["f1"]}

eval_pred 是 Trainer 在每次评估后传入的 (logits, labels) 元组。np.argmax 选择概率（logit）最高的标签作为预测结果。

8. 初始化 Trainer 并训练

from transformers import Trainer, DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()

Epoch   Training Loss   Validation Loss   Accuracy   F1
1       0.693100        0.671204          0.500000   0.500000
2       0.654200        0.598311          1.000000   1.000000
3       0.581900        0.520447          1.000000   1.000000

DataCollatorWithPadding 按批次处理动态填充（比 tokenize 时的固定 max_length 填充更高效，但这里已经预填充，collator 主要负责创建批次张量）。trainer.train() 执行前向传播 → 反向传播 → 优化器步骤 → 记录损失的循环，并根据配置的 eval_strategy 自动评估。

8.1. 监控训练过程

如果想将详细日志输出到 TensorBoard，在 TrainingArguments 中添加 report_to="tensorboard"，然后运行 tensorboard --logdir ./order-review-model/runs。对于演示这样的小数据集，控制台日志就足以检查损失是否在逐渐减少。

9. 评估和保存模型

metrics = trainer.evaluate()
print(metrics)

trainer.save_model("./order-review-model/final")
tokenizer.save_pretrained("./order-review-model/final")

{'eval_loss': 0.5204, 'eval_accuracy': 1.0, 'eval_f1': 1.0, 'eval_runtime': 0.12, ...}

trainer.evaluate() 在 eval_dataset 上重新运行并返回度量字典。save_model 保存配置和权重；记住也要 save_pretrained tokenizer，以便重新加载时词汇表不会出错。

10. 使用微调模型进行推理

使用 pipeline 快速测试刚训练好的模型。

from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="./order-review-model/final",
    tokenizer="./order-review-model/final",
)

result = classifier("配送准时，产品比预期更好看")
print(result)

[{'label': 'positive', 'score': 0.9123}]

pipeline 自动处理 tokenize、前向传播和 softmax，返回带有 score（置信度概率）的 label。这是将模型集成到小型服务中的最快方式，例如 PaymentService 需要自动分类客户反馈。

11. 编写模型行为测试

对于演示数据集，我们可以使用 pytest 编写简单的测试，确保 pipeline 正常工作，更改 checkpoint 时不会出现回归。

import pytest
from transformers import pipeline

@pytest.fixture(scope="module")
def classifier():
    return pipeline(
        "text-classification",
        model="./order-review-model/final",
        tokenizer="./order-review-model/final",
    )

def test_positive_review_is_classified_correctly(classifier):
    result = classifier("产品很棒，配送迅速")[0]

    assert result["label"] == "positive"
    assert result["score"] > 0.5

def test_negative_review_is_classified_correctly(classifier):
    result = classifier("产品有缺陷，配送迟到，非常失望")[0]

    assert result["label"] == "negative"
    assert result["score"] > 0.5

scope="module" 使 pipeline 在文件中所有测试中只加载一次模型，避免每个测试用例都重新加载消耗时间。这个测试只检查高级行为（标签方向正确），不应该用来评估模型质量—那是上面评估步骤的工作。

12. 常见错误

• 忘记 padding/truncation: 超过 max_length 的句子会导致运行时错误，因为批次中的张量大小不一致。
• 搞混 num_labels: 如果数据集有 3 个类别但声明 num_labels=2，Trainer 仍会运行但结果完全错误，因为分类头的输出不足。
• 数据集太小却期望在真实数据上取得高 F1: 例如本文中只是为了演示代码流程，4 行训练数据不足以让模型学到一般规律。不要以这里的 accuracy=1.0 作为标准。
• 没有设置 seed: 每次运行的训练结果都会不同，难以在调试超参数时比较。在训练前使用 transformers 中的 set_seed(42)。
• 在低配机器上使用大型模型（base/large）: 应该从 distilbert/tiny 版本开始，加快开发循环，只有在 pipeline 确定后才切换到大型模型。

何时不应该使用 Trainer: 如果你需要完全自定义训练循环（例如复杂的多任务损失、难以通过回调表达的自定义调度器），直接编写纯 PyTorch 循环或使用 accelerate 会比强行融入 Trainer 更灵活。

13. 总结

Trainer API 帮助你微调小型模型，无需自己编写训练/评估循环，只需准备数据集、tokenizer、模型和 TrainingArguments。对于订单评论分类这样的简单任务，标准流程是：tokenize → 配置 TrainingArguments → 定义 compute_metrics → 初始化 Trainer → train() → evaluate() → save_model()。当处理更复杂的真实数据时，你只需更换数据集和 checkpoint，pipeline 的其他部分基本保持不变。