PyTorch 自然语言处理：Transformers 与微调

FreeGuideOnline 最新 2026-06-12

PyTorch 自然语言处理实战：Transformers 与微调

本教程将带你从零开始，使用 PyTorch 和 Hugging Face 的 Transformers 库，完成一个真实的自然语言处理（NLP）任务——文本分类微调。即使你刚刚接触深度学习，只要有一些 Python 基础，就能跟着本教程一步步把预训练语言模型变成你专属的分类利器。我们将使用经典的 BERT 模型，对电影评论数据集进行情感二分类（正面/负面）。

你将学到什么

  • 搭建 PyTorch 与 Transformers 工作环境
  • 理解预训练模型与微调的基本思想
  • 加载并预处理文本数据集
  • 使用 Trainer API 进行高效微调
  • 评估模型并在新数据上进行预测
  • 保存并复用微调后的模型

环境准备

首先确保你的 Python 版本 ≥ 3.8。建议在虚拟环境中操作。以下命令将安装所需的核心库：

pip install torch transformers datasets evaluate accelerate
  • torch：PyTorch 深度学习框架
  • transformers：Hugging Face 预训练模型库
  • datasets：快速加载与处理数据集
  • evaluate：便捷的模型评估工具
  • accelerate：让训练过程更高效，无需复杂配置

安装完成后，打开 Python 并验证 GPU 是否可用（如果使用 CPU 训练也可以，但会更耗时）：

import torch
print(torch.cuda.is_available())  # True 表示可以使用 GPU

在本教程中，我们默认使用 CPU 也可以完成训练，但推荐使用 GPU 环境（如 Google Colab 免费 GPU）以获得流畅体验。

理解预训练与微调

现代 NLP 的成功很大程度上源于预训练 + 微调范式：

  1. 预训练模型：在大规模语料（如整个维基百科）上训练好的通用语言模型，例如 BERT、RoBERTa 等。它们已经学会了丰富的语言知识。
  2. 微调（Fine-tuning）：在你的特定任务数据上继续训练该模型，让它适配你的任务。由于起点很高，通常只需要很少的数据和训练步数就能得到出色结果。

我们选择 bert-base-uncased 作为基础模型，它拥有 12 层 Transformer，约 1.1 亿参数，输入文本将全部转为小写处理。

加载数据集

我们将使用 IMDb 电影评论数据集，它包含 25000 条训练样本和 25000 条测试样本，标签为正面（1）或负面（0）。Hugging Face 的 datasets 库可以直接加载：

from datasets import load_dataset

dataset = load_dataset("imdb")
print(dataset)

输出大致如下：

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    unsupervised: Dataset({
        features: ['text', 'label'],
        num_rows: 50000
    })
})

我们只使用 train 和 test。为了加速示范，我们取训练集的一个子集（比如 2000 条）和测试集的 1000 条。你可以通过切片轻松实现：

small_train = dataset["train"].shuffle(seed=42).select(range(2000))
small_test = dataset["test"].shuffle(seed=42).select(range(1000))

文本预处理：分词与编码

预训练模型不能直接理解原始文本，需要先将其转换为模型能处理的数字序列。这个过程叫分词（Tokenization），由模型配套的分词器完成。

加载分词器：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

定义一个预处理函数，让分词器对文本进行截断和填充：

def preprocess_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,       # 超过最大长度的文本将被截断
        padding="max_length",  # 填充到固定长度
        max_length=256         # BERT 最大长度为 512，这里设为 256 以加快训练
    )

批量应用到数据集上：

encoded_train = small_train.map(preprocess_function, batched=True)
encoded_test = small_test.map(preprocess_function, batched=True)

现在数据集中多了三个重要字段：

  • input_ids：词编码后的数字序列
  • attention_mask：指示哪些位置是真实 token（1）哪些是填充（0）
  • token_type_ids：用于句子对任务，本文用不到但依然存在

我们还需要确保 PyTorch 能正确读取这些字段，因此设置数据集格式：

encoded_train.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
encoded_test.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])

加载预训练模型

现在载入用于序列分类的 BERT 模型。AutoModelForSequenceClassification 会根据模型名称自动加载合适的结构，并添加一个分类头。

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased",
    num_labels=2  # 二分类，所以输出维度为 2
)

输出中的警告可以忽略。模型权重随机初始化了分类头，其余部分使用了预训练参数。

定义评估指标

我们需要在训练过程中监控模型的准确率。借助 evaluate 库可以轻松加载常用指标：

import evaluate
import numpy as np

accuracy = evaluate.load("accuracy")

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

compute_metrics 函数会在每个评估阶段被调用，输入是模型预测的 logits 和真实标签。

设置训练参数并使用 Trainer 微调

Transformers 提供了高级的 Trainer 类，将训练循环、日志、评估、保存等功能高度封装。我们只需要定义训练参数：

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./results",          # 模型保存路径
    num_train_epochs=3,              # 训练轮数
    per_device_train_batch_size=16,  # 训练批次大小
    per_device_eval_batch_size=16,   # 评估批次大小
    eval_strategy="epoch",           # 每个 epoch 结束进行一次评估
    save_strategy="epoch",           # 每个 epoch 保存一次模型
    logging_dir="./logs",            # 日志目录
    logging_steps=10,
    load_best_model_at_end=True,     # 训练结束后加载表现最好的模型
    metric_for_best_model="accuracy",# 用于选择最佳模型的指标
)

然后创建 Trainer 实例并开始训练：

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=encoded_train,
    eval_dataset=encoded_test,
    compute_metrics=compute_metrics,
)

trainer.train()

训练过程中你将看到类似以下的进度条和损失下降：

Epoch 1/3: 100%|██████████| ... loss=0.3547 ...
Epoch 2/3: 100%|██████████| ... loss=0.1284 ...
Epoch 3/3: 100%|██████████| ... loss=0.0459 ...

由于我们只用了 2000 条数据，准确率可能在 80% 左右，但足以说明微调的有效性。如果你使用全量数据，准确率很容易超过 90%。

评估模型

训练结束后，我们可以直接在测试集上评估最终模型：

results = trainer.evaluate()
print(results)

输出会包含 eval_loss 和 eval_accuracy。你也可以使用 trainer.predict(encoded_test) 获取详细预测概率。

进行新样本推理

微调是为了让模型能够对未见过的文本做出判断。现在它已经学会了情感分类，我们来试试几条自定义评论：

import torch

def predict_sentiment(text):
    # 分词并转换为 PyTorch 张量
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=256)
    # 不计算梯度，加快推理
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    probabilities = torch.nn.functional.softmax(logits, dim=-1)
    pred = torch.argmax(probabilities, dim=-1).item()
    return "正面" if pred == 1 else "负面", probabilities[0].tolist()

test_comments = [
    "This movie was fantastic! I really enjoyed every minute of it.",
    "Waste of time, terrible acting and boring plot."
]

for comment in test_comments:
    sentiment, probs = predict_sentiment(comment)
    print(f"评论: {comment}\n预测: {sentiment} (正: {probs[1]:.4f}, 负: {probs[0]:.4f})\n")

你应该会看到第一个评论被预测为“正面”且概率很高，第二个评论被预测为“负面”。

保存与复用模型

将微调后的模型和分词器保存到本地目录：

model.save_pretrained("./my_finetuned_bert")
tokenizer.save_pretrained("./my_finetuned_bert")

以后使用时，只需一行代码即可加载：

from transformers import pipeline

classifier = pipeline("text-classification", model="./my_finetuned_bert", tokenizer="./my_finetuned_bert")
print(classifier("I loved the special effects in this film."))

你也可以将模型推送到 Hugging Face Hub 与全球开发者分享（需要注册账号并登录）。

进阶方向

恭喜你完成了第一个 PyTorch 自然语言处理微调实战！从这里出发，你还可以尝试：

  • 切换不同预训练模型：例如 roberta-base、distilbert-base-uncased（更小更快）或 albert-base-v2
  • 处理其他任务：使用 AutoModelForTokenClassification 进行命名实体识别，AutoModelForQuestionAnswering 做问答系统
  • 调整超参数：学习率、批次大小、训练轮数对效果影响显著
  • 自定义数据集：加载 CSV 或 JSON 文件，替换 IMDb 数据集
  • 早停和回调：利用 EarlyStoppingCallback 防止过拟合

常见问题

Q: 训练时内存不足怎么办？
A: 减小 per_device_train_batch_size 或 max_length，也可以使用梯度累积（设置 gradient_accumulation_steps）。

Q: 能否在 CPU 上训练？
A: 可以，但速度会慢很多。对于小数据集（如本例）依然可行，只需耐心等待即可。

Q: 微调后的模型只对电影评论有效？
A: 是的，模型学到了你提供的训练数据的分布。如果希望适用于电商评论等场景，需要用相应领域的数据进行微调。

总结

本教程通过一个端到端的示例，展示了如何使用 PyTorch 和 Transformers 对 BERT 进行文本分类微调。你学会了加载数据、分词、定义训练流程、评估和推理的全过程。这个框架可以无缝迁移到大多数 NLP 任务中，为你的智能应用提供强大的语言理解能力。

现在，带着你的新技能，去创造属于自己的 NLP 应用吧！