Hugging Face transformers 微调 BERT 做文本分类（最小流程）

预训练 + 微调是 NLP 的标准范式。Hugging Face transformers 把
模型 / tokenizer / 训练循环都封装好，从原始数据到训练好的分类器
只需 < 50 行代码。

下面用 IMDB 影评数据微调 BERT 做二分类（正面 / 负面）。

装

uv add 'transformers[torch]' datasets accelerate evaluate

数据

from datasets import load_dataset

ds = load_dataset('imdb')
print(ds)
# DatasetDict({
#     train: Dataset({ features: ['text', 'label'], num_rows: 25000 })
#     test:  Dataset({ features: ['text', 'label'], num_rows: 25000 })
#     unsupervised: ...
# })

# label: 0 = neg, 1 = pos
print(ds['train'][0])

tokenizer

from transformers import AutoTokenizer

MODEL = 'distilbert-base-uncased'   # 比 bert-base 小 40%，效果差几个点但快
tok = AutoTokenizer.from_pretrained(MODEL)

def tokenize(batch):
    return tok(batch['text'], truncation=True, max_length=256, padding='max_length')

ds_tok = ds.map(tokenize, batched=True)
ds_tok = ds_tok.remove_columns(['text'])
ds_tok = ds_tok.rename_column('label', 'labels')
ds_tok.set_format('torch')

truncation=True 截到 256 token；BERT 最大 512，但短点训练快得多。

模型

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2)

from_pretrained 自动下载预训练权重 + 加上一个新的 classification head。

Trainer：30 行训练循环

from transformers import TrainingArguments, Trainer
import evaluate
import numpy as np

metric = evaluate.load('accuracy')
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=1)
    return metric.compute(predictions=preds, references=labels)

args = TrainingArguments(
    output_dir='./out',
    num_train_epochs=2,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    weight_decay=0.01,
    warmup_steps=500,
    logging_steps=50,
    eval_strategy='epoch',
    save_strategy='epoch',
    load_best_model_at_end=True,
    metric_for_best_model='accuracy',
    fp16=True,                   # GPU 时开混合精度
    report_to='wandb',           # 可选：wandb 跟踪
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=ds_tok['train'].shuffle(seed=42).select(range(10000)),
    eval_dataset=ds_tok['test'].select(range(2000)),
    tokenizer=tok,
    compute_metrics=compute_metrics,
)

trainer.train()
trainer.evaluate()

10k 训练样本 / 2k 验证样本，DistilBERT，单个 RTX 3090 上 ~5 分钟训完，
准确率约 91-92%。

全量 25k 训 3 个 epoch 能到 93-94%。

保存 + 加载 + 推理

trainer.save_model('./imdb-distilbert')
tok.save_pretrained('./imdb-distilbert')

# 加载推理
from transformers import pipeline
classifier = pipeline('sentiment-analysis', model='./imdb-distilbert',
                       device=0)   # device=0 = cuda:0; -1 = cpu

print(classifier('This movie was absolutely fantastic'))
# [{'label': 'LABEL_1', 'score': 0.998}]

print(classifier('What a complete waste of time and money'))
# [{'label': 'LABEL_0', 'score': 0.997}]

label 是模型内部的，可以在训练时设标签名：

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL, num_labels=2,
    id2label={0: 'NEGATIVE', 1: 'POSITIVE'},
    label2id={'NEGATIVE': 0, 'POSITIVE': 1},
)

之后 pipeline 直接返回 'POSITIVE' / 'NEGATIVE'。

中文文本？

MODEL = 'bert-base-chinese'                    # 谷歌中文 BERT
# 或：
MODEL = 'hfl/chinese-roberta-wwm-ext'          # 哈工大 RoBERTa 全词 mask
# 或：
MODEL = 'IDEA-CCNL/Erlangshen-Roberta-110M-Sentiment'  # 已经是情感分类微调过的

中文 tokenizer 是字级 BPE（不像英文按 subword），输出结构一样。

推到 Hugging Face Hub（可选）

huggingface-cli login

trainer.push_to_hub('your-username/imdb-distilbert')
# 之后任何人能：
# AutoModelForSequenceClassification.from_pretrained('your-username/imdb-distilbert')

减小内存：LoRA

完整微调 DistilBERT 已经够轻；微调 7B+ 模型时显存装不下，用 LoRA：

uv add peft

from peft import LoraConfig, get_peft_model, TaskType

lora_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    r=8,
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules=['q_lin', 'v_lin'],
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 0.6M / 67M = 0.9%

之后照常 Trainer。LoRA 让显存 / 速度大幅降低，性能差距通常 < 1%。

推理优化（生产）

量化：bitsandbytes int8 / int4，2-4x 速度，半精度差几个点
ONNX 导出 + ONNX Runtime：CPU 推理 2-3x 加速
TGI (Text Generation Inference) / vLLM：高并发 LLM 推理
TorchServe：通用 PyTorch 推理服务

简单分类任务直接用 pipeline；高并发用 TGI / TorchServe。

评估更精细

import evaluate

metrics = evaluate.combine(['accuracy', 'f1', 'precision', 'recall'])

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

踩过的坑

padding='max_length' 把所有样本 pad 到 256 → 浪费计算。改成
padding=True（动态 pad 到 batch 内最长），训练快 2-3x。
learning rate 太大：BERT 微调通常 2e-5 ~ 5e-5。1e-4 已经偏大，
常导致 loss NaN。
数据集 label 顺序：HuggingFace load_dataset('imdb') 是按 label
排序的，不 shuffle 训练会先看完所有 negative 再看 positive，
loss 曲线奇形怪状。shuffle(seed=42) 必加。
fp16=True 在 BF16 友好的硬件（A100 / H100）改 bf16=True 更稳。
老 V100 / GTX 用 fp16。