文本分类实战：基于FastText和BERT的完整代码实现与避坑指南

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

1. 技术选型：为什么选择 FastText 和 BERT？

在文本分类任务中，模型的选择直接影响最终效果和实现复杂度。对于初学者而言，避开过于简单的朴素贝叶斯和需要复杂调参的 TextCNN，我们推荐以下两种方案：

• FastText：Facebook 开源的轻量级模型，特别适合小规模数据集（10 万条以下）。其核心优势在于：
• 通过 n -gram 特征捕捉局部词序信息
• 内置词向量训练和层次 softmax 加速
• 训练速度快（CPU 即可运行）

• 代码实现仅需 5 -10 行

• BERT：Google 提出的预训练模型，适合复杂场景（如长文本、多语义）：

• 基于 Transformer 的深度双向编码
• 自动学习上下文相关特征
• 微调即可达到 SOTA 效果
• 需要 GPU 支持

2. 完整代码实现

2.1 数据预处理

无论使用哪种模型，都需要先完成以下预处理步骤（以中文新闻分类为例）：

import jieba
import pandas as pd
from sklearn.model_selection import train_test_split

# 示例数据加载
data = pd.read_csv('thucnews.csv')  # 假设列名为 text, label

# 文本清洗
def clean_text(text):
    text = re.sub(r'[\s\n\r\t]', '', text)  # 去除空白符
    text = re.sub(r'[^\w\u4e00-\u9fff]', '', text)  # 保留中英文和数字
    return text

# 中文分词
def chinese_seg(text):
    return ' '.join(jieba.cut(clean_text(text)))

data['processed'] = data['text'].apply(chinese_seg)

# 划分训练测试集
X_train, X_test, y_train, y_test = train_test_split(data['processed'], data['label'], test_size=0.2)

2.2 FastText 实现

安装 FastText 库后，模型训练仅需几行代码：

import fasttext

# 准备 FastText 要求的文件格式
with open('train.txt', 'w') as f:
    for text, label in zip(X_train, y_train):
        f.write(f'__label__{label} {text}\n')  # 标签需加__label__前缀

# 模型训练（关键参数说明）model = fasttext.train_supervised(
    input='train.txt',
    lr=0.1,           # 学习率
    dim=100,          # 词向量维度
    epoch=50,         # 迭代次数
    wordNgrams=2,     # 使用 2 -gram 特征
    loss='hs'         # 层次 softmax 加速训练
)

# 评估
print(model.test('test.txt'))  # 准确率与样本数

2.3 BERT 微调方案

使用 HuggingFace Transformers 库实现 BERT 微调：

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

# 加载预训练模型和分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertForSequenceClassification.from_pretrained(
    'bert-base-chinese', 
    num_labels=len(y_train.unique())
)

# 构建 Dataset 类
class NewsDataset(Dataset):
    def __init__(self, texts, labels):
        self.encodings = tokenizer(texts.tolist(), truncation=True, padding=True)
        self.labels = labels.tolist()

    def __getitem__(self, idx):
        item = {k: torch.tensor(v[idx]) for k, v in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

# 训练循环（简化版）train_loader = DataLoader(NewsDataset(X_train, y_train), batch_size=16)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)

for epoch in range(3):  # BERT 微调通常 3 - 5 个 epoch 足够
    for batch in train_loader:
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()
        optimizer.step()

3. 避坑指南

3.1 类别不平衡处理

• FastText：使用 -weight 参数设置类别权重

  model = fasttext.train_supervised(
      input='train.txt',
      weight='1,2,3'  # 根据类别频率设置权重
  )

• BERT：在损失函数中加权

  from torch.nn import CrossEntropyLoss
  loss_fct = CrossEntropyLoss(weight=torch.tensor([1.0, 2.0, 0.5]))

3.2 部署优化

• 内存优化：
• FastText：.quantize()方法可压缩模型（牺牲少量精度）

• BERT：使用 bert-mini 等轻量变体

• 延迟控制：

• FastText：单条预测约 1ms（CPU）
• BERT：通过 ONNX 转换 +TensorRT 加速

4. 性能对比

在 THUCNews 数据集（10 类别）上的测试结果：

5. 总结与思考

根据业务场景选择模型的决策路径：

1. 数据量少（<10 万条）→ FastText
2. 需要最高精度 → BERT
3. 无 GPU 资源 → FastText
4. 需要实时预测 → FastText

扩展实验建议：
– 尝试 FastText 与 BERT 的集成（如用 BERT 特征作为 FastText 输入）
– 测试不同语言的性能差异
– 探索知识蒸馏压缩 BERT 模型