BERT-base-chinese模型多分类微调实战：从数据准备到生产部署

1次阅读

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

中文文本分类相比英文面临更多挑战。首先，中文没有明显的单词分隔符，分词效果直接影响模型性能。其次，中文的字符组合更加灵活，同一个词在不同上下文可能有不同含义。另外，中文数据标注成本高，很多领域存在数据稀疏问题。这些因素都使得直接使用预训练模型往往无法达到理想效果，微调变得尤为重要。

在中文 NLP 任务中，常见的预训练模型主要有以下几种：

BERT-base-chinese：专门针对中文优化的 BERT 模型，在大多数中文任务上表现稳定
RoBERTa-wwm-ext：使用全词掩码技术，对中文实体识别效果更好
ALBERT：参数共享机制减少模型大小，适合资源受限场景

经过对比测试，我们发现对于大多数多分类任务，BERT-base-chinese 在准确率和推理速度上取得了较好的平衡。它的 12 层 Transformer 结构足够捕捉中文语义，同时模型大小 (110M 参数) 也适合生产部署。

中文文本预处理有几个关键点需要注意：

统一字符编码：确保所有文本使用 UTF- 8 编码
特殊字符处理：过滤或替换网页标签、表情符号等非常规字符
文本长度控制：BERT 最大长度限制为 512，需要合理截断或填充

以下是预处理代码示例：

import re
from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')

def clean_text(text):
    # 移除 HTML 标签
    text = re.sub(r'<[^>]+>', '', text)
    # 替换特殊空白字符
    text = re.sub(r'[\u3000\xa0]', ' ', text)
    # 过滤非中文字符和标点
    text = re.sub(r'[^\w\s\p{P}\p{Han}]', '', text)
    return text.strip()

def tokenize_with_bert(text, max_length=128):
    text = clean_text(text)
    return tokenizer(
        text, 
        max_length=max_length, 
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )

BERT 默认的输出层不适合多分类任务，我们需要修改分类头：

import torch.nn as nn
from transformers import BertModel

class BertForMultiClass(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.bert = BertModel.from_pretrained('bert-base-chinese')
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(768, num_classes)  # BERT 隐藏层维度是 768

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        pooled_output = outputs.pooler_output
        pooled_output = self.dropout(pooled_output)
        return self.classifier(pooled_output)

中文数据常常存在类别不平衡问题，我们使用 Focal Loss 替代标准交叉熵：

class FocalLoss(nn.Module):
    def __init__(self, alpha=1, gamma=2, reduction='mean'):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction

    def forward(self, inputs, targets):
        bce_loss = nn.functional.cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-bce_loss)
        focal_loss = self.alpha * (1-pt)**self.gamma * bce_loss

        if self.reduction == 'mean':
            return focal_loss.mean()
        elif self.reduction == 'sum':
            return focal_loss.sum()
        return focal_loss

使用 AMP(Automatic Mixed Precision)可以显著减少显存占用并加快训练速度：

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

for batch in train_loader:
    optimizer.zero_grad()

    with autocast():
        outputs = model(batch['input_ids'], batch['attention_mask'])
        loss = criterion(outputs, batch['labels'])

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

BERT 微调时，前几个 epoch 使用较小的学习率有助于稳定训练：

from transformers import get_linear_schedule_with_warmup

num_training_steps = len(train_loader) * epochs
num_warmup_steps = num_training_steps // 5

scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=num_warmup_steps,
    num_training_steps=num_training_steps
)

监控验证集损失，当连续若干 epoch 没有改善时停止训练：

best_loss = float('inf')
patience = 3
no_improve = 0

for epoch in range(epochs):
    # 训练和验证代码...

    if val_loss < best_loss:
        best_loss = val_loss
        no_improve = 0
        torch.save(model.state_dict(), 'best_model.pt')
    else:
        no_improve += 1
        if no_improve >= patience:
            print(f'Early stopping at epoch {epoch}')
            break

使用梯度累积：每 N 个 batch 更新一次参数
减小 batch size
启用梯度检查点：model.gradient_checkpointing_enable()

BERT 的 tokenizer 首次运行较慢，可以预先缓存：

# 初次使用前先缓存
for text in tqdm(texts):
    tokenizer(text)

# 保存 tokenizer 缓存
tokenizer.save_pretrained('./tokenizer_cache/')

使用动态量化减小模型体积，提升推理速度：

import torch.quantization

quantized_model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8
)
torch.save(quantized_model.state_dict(), 'quantized_model.pt')

我们在三个不同规模的中文数据集上进行了测试：