基于Transformer的文本分类实战：从数据预处理到模型部署完整指南

共计 2378 个字符，预计需要花费 6 分钟才能阅读完成。

问题背景

传统文本分类模型如朴素贝叶斯和 TextCNN 在处理复杂语义时存在明显短板：

• 朴素贝叶斯无法捕捉词序信息，且依赖强独立性假设
• TextCNN 虽然能处理局部特征，但对长距离依赖建模能力有限
• RNN 系列模型存在梯度消失问题，难以处理超长文本

Transformer 架构通过自注意力机制解决了这些问题，成为当前 NLP 任务的黄金标准。

技术选型

本次实战选择 HuggingFace 生态，核心优势包括：

1. 丰富的预训练模型库（BERT/RoBERTa/DistilBERT 等）
2. 统一的 Pipeline 接口
3. 完善的社区支持

实验环境要求：
– Python 3.8+
– PyTorch 1.12+
– CUDA 11.3
– transformers 4.25+

实现细节

数据准备模块

from sklearn.preprocessing import LabelEncoder
from torch.utils.data import Dataset
import pandas as pd
import re

class TextDataset(Dataset):
    """自定义数据集类"""
    def __init__(self, texts: list[str], labels: list[int], tokenizer, max_len: int = 512):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self) -> int:
        return len(self.texts)

    def __getitem__(self, idx: int) -> dict:
        text = str(self.texts[idx])
        # 数据清洗
        text = re.sub(r'[^\w\s]', '', text)  # 移除特殊字符
        encoding = self.tokenizer(
            text,
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'label': torch.tensor(self.labels[idx], dtype=torch.long)
        }

模型训练核心逻辑

import torch
from transformers import BertForSequenceClassification, AdamW
from torch.optim.lr_scheduler import ReduceLROnPlateau

# 初始化模型
model = BertForSequenceClassification.from_pretrained(
    'bert-base-uncased',
    num_labels=num_classes
)

# 带权重衰减的优化器
optimizer = AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)

# 动态学习率调整
scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.1, patience=3)

# 类别不平衡处理
weights = torch.tensor([1.0, 2.0, 3.0])  # 假设第二类比第一类重要 2 倍
criterion = torch.nn.CrossEntropyLoss(weight=weights)

性能优化

处理类别不平衡的三种策略

1. 重采样技术
2. 过采样少数类（SMOTE）

3. 欠采样多数类

4. Loss 函数调整

5. Focal Loss

6. 带权重的 CrossEntropy

7. 评估指标优化

8. 改用 Macro-F1 替代 Accuracy

模型轻量化方案

# 知识蒸馏示例
from transformers import DistilBertForSequenceClassification

distilled_model = DistilBertForSequenceClassification.from_pretrained(
    'distilbert-base-uncased',
    num_labels=num_classes
)

# 量化推理
quantized_model = torch.quantization.quantize_dynamic(
    model,
    {torch.nn.Linear},
    dtype=torch.qint8
)

生产实践

模型部署流水线

1. ONNX 格式转换

   torch.onnx.export(
       model,
       dummy_input,
       "bert_text_cls.onnx",
       opset_version=13,
       input_names=['input_ids', 'attention_mask'],
       output_names=['logits']
   )

2. Triton 推理服务配置

   platform: "onnxruntime_onnx"
   max_batch_size: 32
   input [
     {
       name: "input_ids"
       data_type: TYPE_INT64
       dims: [-1, 512]
     }
   ]

总结展望

通过本实践我们实现了：
– 基于 BERT 的端到端文本分类流程
– 工业级优化技巧
– 生产环境部署方案

延伸思考方向：
1. 长文本处理可采用：
– Longformer 的稀疏注意力
– 文本分块 + 投票策略
2. 参数高效微调：
– LoRA 仅训练 0.1% 参数达到 90% 全量微调效果
– Adapter 引入少量可训练层

完整项目代码已开源：github.com/example/text-classification-transformers