BERT微调实战：如何在自己的数据集上高效训练模型

1次阅读

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

BERT（Bidirectional Encoder Representations from Transformers）是一种预训练语言模型，通过微调可以适应各种下游 NLP 任务。微调是指在预训练模型的基础上，使用特定领域的数据进行额外训练，使模型更好地适应目标任务。

文本分类（情感分析、新闻分类等）
命名实体识别
问答系统
文本相似度计算

BERT 要求输入数据具有特定格式，而原始数据往往不符合要求。例如：

文本长度不一致
标签格式不符合模型输出
特殊字符未处理

BERT 模型参数量大，训练需要大量 GPU 资源，尤其是当数据集较大时。

在小数据集上微调 BERT 容易导致过拟合，模型在训练集上表现良好但在测试集上表现不佳。

文本清洗：去除特殊字符、HTML 标签等
分词：使用 BERT 的 Tokenizer
生成输入 ID 和注意力掩码
划分训练集、验证集和测试集

加载预训练 BERT 模型
定义自定义分类头
设置优化器和学习率调度器
训练循环
模型评估

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

# 1. 自定义数据集类
class CustomDataset(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',
            return_attention_mask=True,
            return_tensors='pt',
            truncation=True
        )

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

# 2. 数据准备
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# 假设我们有以下数据
texts = ["This is a positive example", "This is negative"]
labels = [1, 0]

# 3. 创建数据集和数据加载器
dataset = CustomDataset(texts, labels, tokenizer, max_len=128)
data_loader = DataLoader(dataset, batch_size=16, shuffle=True)

# 4. 模型初始化
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
model = model.to('cuda' if torch.cuda.is_available() else 'cpu')

# 5. 训练设置
optimizer = AdamW(model.parameters(), lr=2e-5)
epochs = 3

# 6. 训练循环
for epoch in range(epochs):
    model.train()
    for batch in data_loader:
        optimizer.zero_grad()

        input_ids = batch['input_ids'].to(model.device)
        attention_mask = batch['attention_mask'].to(model.device)
        labels = batch['label'].to(model.device)

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

        loss = outputs.loss
        loss.backward()
        optimizer.step()

    print(f'Epoch {epoch+1}/{epochs}, Loss: {loss.item()}')

BERT 微调通常使用较小的学习率（2e- 5 到 5e-5），可以使用学习率预热策略：

from transformers import get_linear_schedule_with_warmup

# 在训练前添加
num_training_steps = len(data_loader) * epochs
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=0,
    num_training_steps=num_training_steps
)

根据 GPU 内存选择合适的批次大小，通常 16-32 效果较好。如果内存不足，可以使用梯度累积：

# 修改训练循环
accumulation_steps = 4
for step, batch in enumerate(data_loader):
    # 前向传播和计算损失
    loss = loss / accumulation_steps
    loss.backward()

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