BERT模型微调实战：从零开始构建高效NLP任务的代码指南

1次阅读

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

BERT（Bidirectional Encoder Representations from Transformers）是 Google 在 2018 年提出的预训练语言模型，通过大规模无监督学习获得了强大的语言表示能力。微调（Fine-tuning）是指在特定任务上对预训练好的 BERT 模型进行少量训练，使其适应新的任务需求。

BERT 微调常见的应用场景包括：

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

初学者在 BERT 微调过程中常遇到以下问题：

数据格式不匹配：BERT 需要特定的输入格式（如 token IDs、attention masks 等）
训练效率低：没有合理设置 batch size 和学习率等超参数
内存不足：BERT 模型较大，容易导致显存溢出
过拟合：在小数据集上微调时容易出现

以下是使用 PyTorch 和 HuggingFace Transformers 库进行 BERT 微调的完整代码示例：

# 导入必要库
import torch
from transformers import BertTokenizer, BertForSequenceClassification
from transformers import AdamW, get_linear_schedule_with_warmup
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader, Dataset

# 1. 数据准备
class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

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

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

        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(),
            'labels': torch.tensor(label, dtype=torch.long)
        }

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

# 假设我们有一些文本数据和标签
texts = ["I love this movie", "This movie is terrible", ...]
labels = [1, 0, ...]  # 1 表示正面，0 表示负面

# 分割训练集和测试集
train_texts, val_texts, train_labels, val_labels = train_test_split(texts, labels, test_size=0.1, random_state=42)

# 创建数据加载器
train_dataset = TextDataset(train_texts, train_labels, tokenizer)
val_dataset = TextDataset(val_texts, val_labels, tokenizer)

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=16)

# 3. 设置优化器和学习率调度器
optimizer = AdamW(model.parameters(), lr=2e-5, correct_bias=False)
total_steps = len(train_loader) * 3  # 假设训练 3 个 epoch

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

# 4. 训练循环
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

for epoch in range(3):  # 训练 3 个 epoch
    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['labels'].to(device)

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

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

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

    avg_train_loss = total_loss / len(train_loader)
    print(f'Epoch {epoch + 1}, Train Loss: {avg_train_loss:.4f}')

    # 验证
    model.eval()
    val_loss = 0
    correct_predictions = 0

    with torch.no_grad():
        for batch in val_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

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

            val_loss += outputs.loss.item()

            _, preds = torch.max(outputs.logits, dim=1)
            correct_predictions += torch.sum(preds == labels)

    avg_val_loss = val_loss / len(val_loader)
    val_acc = correct_predictions.double() / len(val_dataset)
    print(f'Val Loss: {avg_val_loss:.4f}, Val Acc: {val_acc:.4f}')