BERT 文本分类实战：代码逐行注释与原理详解

3. 数据加载与预处理

IMDB 数据集包含 50,000 条评论，其中一半为正面评价，一半为负面评价。我们将首先加载数据并将其转换为数值标签。

3.1 加载函数

def load_imdb_data(data_file):
    """
    加载 IMDB 数据集并转换为文本列表和标签列表
    :param data_file: CSV 文件路径
    :return: texts (列表), labels (列表)
    """
    df = pd.read_csv(data_file)  # 读取 CSV 文件
    texts = df['review'].tolist()  # 提取评论文本
    # 将情感转化为数值：positive -> 1, negative -> 0
    labels = [1 if sentiment == "positive" else 0 for sentiment in df['sentiment'].tolist()]
    return texts, labels

# 指定数据集路径并加载数据
# 注意：实际使用时请确保 'IMDB Dataset.csv' 存在于当前目录
try:
    data_file = "IMDB Dataset.csv"
    texts, labels = load_imdb_data(data_file)
except FileNotFoundError:
    print("未找到数据集文件，请检查路径。")
    texts, labels = [], []

4. 自定义数据集类

PyTorch 的 Dataset 类允许我们自定义数据加载逻辑。我们需要将文本编码为 BERT 可接受的输入格式（Input IDs 和 Attention Mask）。

class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

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

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        
        # 分词与编码
        encoding = self.tokenizer(
            text,
            return_tensors='pt',
            max_length=self.max_length,
            padding='max_length',  # 填充至最大长度
            truncation=True        # 超过长度则截断
        )
        
        return {
            'input_ids': encoding['input_ids'].squeeze(0),  # 移除 batch 维度
            'attention_mask': encoding['attention_mask'].squeeze(0),
            'label': torch.tensor(label)
        }

5. 构建 BERT 分类器

我们需要在预训练的 BERT 模型之上添加一个分类头（Classifier Head）。BERT 的输出经过池化层后，接入全连接层进行二分类。

class BERTClassifier(nn.Module):
    def __init__(self, bert_model_name, num_classes):
        super(BERTClassifier, self).__init__()
        # 加载预训练 BERT 模型
        self.bert = BertModel.from_pretrained(bert_model_name)
        # Dropout 层防止过拟合
        self.dropout = nn.Dropout(0.1)
        # 全连接层：从 BERT 隐藏层大小映射到类别数
        self.fc = nn.Linear(self.bert.config.hidden_size, num_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.pooler_output  # 获取 [CLS] 标记的输出作为句子表示
        x = self.dropout(pooled_output)
        logits = self.fc(x)  # 输出分类得分
        return logits

6. 定义训练与评估函数

6.1 训练函数

训练过程包括前向传播、损失计算、反向传播和优化器更新。

def train(model, data_loader, optimizer, scheduler, device):
    model.train()
    total_loss = 0
    
    for batch in data_loader:
        optimizer.zero_grad()  # 清除上一次的梯度
        
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['label'].to(device)
        
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        loss = nn.CrossEntropyLoss()(outputs, labels)  # 交叉熵损失
        
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数
        scheduler.step()  # 调整学习率
        
        total_loss += loss.item()
    
    return total_loss / len(data_loader)

6.2 评估函数

评估阶段禁用梯度计算以提升效率，并计算准确率及分类报告。

def evaluate(model, data_loader, device):
    model.eval()
    all_labels = []
    all_preds = []
    
    with torch.no_grad():  # 禁用梯度计算
        for batch in data_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['label'].to(device)
            
            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            preds = torch.argmax(outputs, dim=1)  # 获取预测标签
            
            all_labels.extend(labels.cpu().numpy())
            all_preds.extend(preds.cpu().numpy())
    
    acc = accuracy_score(all_labels, all_preds)
    report = classification_report(
        all_labels, all_preds, 
        target_names=["负面", "正面"], 
        digits=4
    )
    return acc, report

7. 超参数设置与模型初始化

合理的超参数选择对模型收敛至关重要。

# 定义参数
bert_model_name = "bert-base-uncased"
max_length = 128  # 文本最大长度，BERT 默认支持 512，但 128 通常足够且更快
batch_size = 16   # 批次大小，取决于 GPU 显存
num_epochs = 4    # 训练轮数
learning_rate = 2e-5  # 学习率，BERT 微调常用值

# 设备选择
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备：{device}")

# 初始化模型
model = BERTClassifier(bert_model_name=bert_model_name, num_classes=2)
model.to(device)

# 优化器和学习率调度器
optimizer = AdamW(model.parameters(), lr=learning_rate, eps=1e-8)
total_steps = len(train_loader) * num_epochs
scheduler = get_linear_schedule_with_warmup(
    optimizer, 
    num_warmup_steps=0, 
    num_training_steps=total_steps
)

8. 数据集划分与加载器创建

将数据划分为训练集和验证集，并使用 DataLoader 进行批量加载。

# 划分训练集和验证集 (80% 训练，20% 验证)
train_texts, val_texts, train_labels, val_labels = train_test_split(
    texts, labels, test_size=0.2, random_state=42
)

# 初始化分词器
tokenizer = BertTokenizer.from_pretrained(bert_model_name)

# 创建自定义数据集
train_dataset = TextClassificationDataset(train_texts, train_labels, tokenizer, max_length)
val_dataset = TextClassificationDataset(val_texts, val_labels, tokenizer, max_length)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)

9. 模型训练流程

执行多轮迭代训练，并在每个 Epoch 结束后验证模型性能。

for epoch in range(num_epochs):
    print(f"\n训练第 {epoch+1}/{num_epochs} 轮...")
    train(model, train_loader, optimizer, scheduler, device)
    acc, report = evaluate(model, val_loader, device)
    
    print(f"验证集准确率：{acc:.4f}")
    print(f"分类报告:\n{report}")
    
    # 可选：保存最佳模型
    # torch.save(model.state_dict(), f"best_model_epoch_{epoch}.pth")

10. 推理与新数据预测

训练完成后，可以使用模型对新数据进行情感预测。

def predict(model, text, tokenizer, max_length, device):
    model.eval()
    
    # 编码输入文本
    encoding = tokenizer(
        text, 
        return_tensors='pt', 
        max_length=max_length,
        padding='max_length', 
        truncation=True
    )
    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    # 获取预测
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        pred = torch.argmax(outputs, dim=1).cpu().numpy()[0]
    
    return "正面" if pred == 1 else "负面"

# 示例预测
example_text = "The movie was absolutely amazing, I loved it!"
prediction = predict(model, example_text, tokenizer, max_length, device)
print(f"评论：{example_text}")
print(f"预测结果：{prediction}")

11. 总结与扩展建议

通过上述步骤，我们成功利用 BERT 实现了高效的文本分类器，并在 IMDB 数据集上完成了情感分析任务。BERT 的强大语言理解能力使得我们无需从零开始训练，大大降低了实现难度。

后续优化方向：

1. 模型保存与加载： 生产环境中应保存 state_dict 以便后续部署。
2. 超参数调优： 尝试不同的学习率、Batch Size 和 Max Length 以获得更高精度。
3. 数据增强： 对于小样本场景，可使用回译等技术扩充训练数据。
4. 分布式训练： 若数据量巨大，可考虑使用多 GPU 并行训练加速。

本教程提供了完整的代码框架与原理讲解，可作为 NLP 入门及 BERT 应用的基础参考。