文本分类实战：BERT与朴素贝叶斯技术选型对比与避坑指南

1次阅读

没有评论

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

文本分类是 NLP 中最基础也最实用的任务之一，比如垃圾邮件过滤、新闻分类、情感分析等场景都会用到。对于初学者来说，最大的挑战往往是面对众多算法时不知道如何选择。今天我们就来对比两种典型的文本分类算法：BERT 和朴素贝叶斯。

朴素贝叶斯基于贝叶斯定理，假设特征之间相互独立。它的核心思想是通过统计词语在不同类别中出现的概率来进行分类。

优点：模型简单、训练速度快、对小数据集友好
缺点：无法捕捉词语间的关联性（因为假设特征独立）

BERT 是基于 Transformer 架构的预训练语言模型，通过自注意力机制理解上下文语义。

优点：能捕捉深层语义、在大型语料上预训练后效果出众
缺点：需要大量计算资源、训练时间长

from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split

# 1. 数据准备
X = [...]  # 文本数据
y = [...]  # 标签

# 2. 特征提取
vectorizer = TfidfVectorizer()
X_vec = vectorizer.fit_transform(X)

# 3. 训练测试分割
X_train, X_test, y_train, y_test = train_test_split(X_vec, y, test_size=0.2)

# 4. 模型训练
model = MultinomialNB()
model.fit(X_train, y_train)

# 5. 评估
print("Accuracy:", model.score(X_test, y_test))

from transformers import BertTokenizer, BertForSequenceClassification
from transformers import Trainer, TrainingArguments
import torch

# 1. 加载预训练模型
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=num_classes)

# 2. 数据预处理
def encode_texts(texts, labels):
    return tokenizer(texts, padding=True, truncation=True, return_tensors="pt")

dataset = encode_texts(X, y)

# 3. 训练参数设置
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=16,
    evaluation_strategy="epoch"
)

# 4. 训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    eval_dataset=dataset
)

trainer.train()

我们在 IMDB 影评数据集上进行了对比测试：