BERT微调实战：从模型选择到生产环境部署的完整指南

1次阅读

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

在实际业务场景中使用 BERT 进行微调时，我们常常遇到以下几个挑战：

计算资源消耗大 ：BERT-base 模型就有 1.1 亿参数，训练需要大量 GPU 内存和算力
小样本学习效果不佳 ：当标注数据有限时，模型容易过拟合
训练效率低下 ：传统的全精度训练速度慢，迭代周期长
生产环境适配困难 ：训练好的模型在部署时可能遇到内存不足、推理延迟高等问题

目前主流的 BERT 微调实现方案主要有三种：

HuggingFace Transformers：
优点：API 设计友好，预训练模型丰富，社区支持好
缺点：部分高级功能需要自行实现
TensorFlow 原生实现 ：
优点：与 TF 生态无缝集成，适合已有 TF 流水线的团队
缺点：代码相对冗长
PyTorch 原生实现 ：
优点：动态图机制调试方便，自定义灵活
缺点：需要手动处理更多底层细节

import torch
from transformers import BertTokenizer, BertForSequenceClassification
from torch.utils.data import Dataset, DataLoader

# 自定义数据集类
class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = str(self.texts[idx])
        label = self.labels[idx]

        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            return_token_type_ids=False,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt'
        )

        return {'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'label': torch.tensor(label, dtype=torch.long)
        }

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

# 准备数据
train_dataset = TextDataset(train_texts, train_labels, tokenizer, max_len=128)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)

# 训练配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
loss_fn = torch.nn.CrossEntropyLoss()

# 训练循环
for epoch in range(3):
    model.train()
    total_loss = 0

    for batch in train_loader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['label'].to(device)

        optimizer.zero_grad()

        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels
        )

        loss = outputs.loss
        total_loss += loss.item()

        loss.backward()
        optimizer.step()

    print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader)}')

学习率调度 ：推荐使用带 warmup 的线性衰减

from transformers import get_linear_schedule_with_warmup

scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=100,
    num_training_steps=len(train_loader)*3
)

早停机制 ：监控验证集 loss，当连续 N 轮不下降时停止训练

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()

for batch in train_loader:
    with autocast():
        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels
        )

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

gradient_accumulation_steps = 4

for step, batch in enumerate(train_loader):
    loss = model(...).loss
    loss = loss / gradient_accumulation_steps
    loss.backward()

    if (step+1) % gradient_accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

使用教师 - 学生模型架构，将大模型知识迁移到小模型：