基于BERT的中文文本分类实战:从模型加载到微调部署 简介基于 pytorch-transformers 的 BERT 中文文本分类完整工程面向有 Python 基础、希望快速落地文本分类任务的开发者与学生覆盖数据预处理、模型微调、评估预测全流程。代码包共 29 个文件包含 5 个 Python 脚本、7 个已划分的文本数据文件、预训练 BERT 权重与词表、结果图表及说明文档等压缩包约 732.57MB。实验数据采用 THUCNews 20 万条新闻标题涵盖财经、房产、股票等 10 个类别按 18 万训练、1 万验证、1 万测试划分类别均衡便于直接复现。目前已有 3699 人学习使用工程结构清晰适合作为课程设计或业务文本分类的参考基线。读者可获得可直接运行的完整代码与模型文件既能复现训练与测试效果也能替换数据目录迁移到自己的数据集并参考可视化结果进行调优。1. 项目概述与整体技术方案1.1 这个项目解决什么问题做中文文本分类的时候很多人第一反应是拿传统机器学习算法去试比如TF-IDF加SVM、逻辑回归这些。说实话如果分类的类别少、数据量不大传统方法确实够用跑起来也快。但一旦遇到语义相近、需要结合上下文才能判断的场景比如“苹果价格走势分析”和“苹果手机评测”这种字面重复但类别完全不同的文本传统方法就很容易翻车。BERT类模型最大的优势在于它用大规模语料预训练出了语义表征能力能够根据上下文动态生成词向量。同一個词放在不同句子里得到的向量表示是不一样。这个特性对中文文本分类特别友好因为中文的语义复杂度高一词多义、歧义现象多。我最初用BERT做中文文本分类就是被这个能力吸引来的实际跑下来效果也确实比之前用的Word2Vec加TextCNN方案高了好几个百分点。这个项目适合两类人看。一类是刚接触NLP、准备用预训练模型做文本分类的开发者可以照着代码把整个流程跑通。另一类是已经用过BERT但想了解中文场景下数据处理细节、训练参数调优和工程化落地注意事项的人。代码本身不复杂核心就是借助pytorch-transformers现在通常叫transformers库加载一个中文预训练模型然后在自己数据集上做微调。1.2 技术选型背后的原因用BERT做分类市面上方案其实不少可以直接用HuggingFace的transformers库也可以自己写模型结构再加载预训练权重甚至可以用封装更上层的paddleNLP或者简单粗暴地调用在线API。我最终选择pytorch-transformers路线有几点具体考量。第一transformers库已经成了事实标准社区活跃遇到问题基本都能搜到解决方案。第二PyTorch生态对研究者更友好调试模型中间层、插入自定义网络结构都方便。第三也是最重要的一点库里面提供了非常方便的加载预训练模型和分词器的方式几行代码就能把模型结构初始化好。对比几种常见中文文本分类方案方案训练速度准确率上限部署难度适用场景TF-IDF 机器学习分类器快中低低长文本、类目差异大的任务Word2Vec TextCNN/BiLSTM较快中中数据量中等、对速度有要求BERT系列微调慢需GPU高中高语义复杂、需要高准确率的任务大模型API调用取决于接口很高低有成本预算充足、不介意数据出域我这里选的是BERT base Chinese版本。它是在中文维基百科等大规模中文语料上预训练出来的词汇表是基于字的对中文支持比较成熟。相比用多语言模型中文BERT在中文任务上的表现通常更稳定。如果数据量大、资金充足也可以换用RoBERTa-wwm-ext、NEZHA这类强化版模型代码基本不用大改只需要换模型名即可。2. 环境准备与数据预处理2.1 环境清单与安装细节先列一下我实际用的环境版本这套组合跑得很稳Python 3.8.10PyTorch 1.10.2CUDA 11.3transformers 4.18.0pandas 1.4.2scikit-learn 1.1.1有一点要提醒pytorch-transformers这个库名在早期版本用过一段时间现在已经统一叫transformers了。如果你搜到一些老教程里写的是from pytorch_transformers import BertModel注意那是旧版API建议直接用新版的from transformers import BertModel接口基本一致但新版本支持更多模型和bug修复。安装命令很简单pip install transformers torch pandas scikit-learn如果你要用GPU训练务必先确认PyTorch的CUDA版本和你本机驱动匹配不然会报“CUDA not available”的错。验证方法python -c import torch; print(torch.cuda.is_available())这一步返回True再往下走。2.2 数据格式规范数据格式是整个项目的地基。BERT模型输入有硬性要求文本必须经过分词器转换成长度固定的token序列。对中文来说BERT的token是把每个字当成一个基本单元典型的ChineseBERT用的是字级别词表不像英文那样把单词切分成子词。这意味着“我喜欢机器学习”会被拆成“我”、“喜”、“欢”、“机”、“器”、“学”、“习”这7个token。我用的数据是CSV格式两列text和label。文本是新闻标题标签是类别编号存档如下text,label 中国男篮备战亚洲杯预选赛,0 苹果发布新一代iPhone手机,1 央行开展逆回购操作维护流动性,2这里有三点实操建议要给到标签务务必转换成整数idBERT分类头输出的是每个类别的logits。文本中不要保留太多特殊符号比如连续的空格、emoji、HTML标签这些会影响分词效果和模型判断。类别数量最好均衡如果某个类别样本太少模型容易偏向多数类后面我会讲怎么缓解。2.3 分词与数据集构建分词这一步直接用transformers库里的BertTokenizer。关键点在于预训练模型都有对应的分词器不要自己写一个分词逻辑去凑必须用和预训练时一致的tokenizer否则token对应的词表id全错了模型效果会断崖式下降。核心代码片段from transformers import BertTokenizer tokenizer BertTokenizer.from_pretrained(bert-base-chinese) def encode_text(text, max_len128): encoded tokenizer.encode_plus( text, max_lengthmax_len, paddingmax_length, truncationTrue, return_tensorspt ) return encodedencode_plus返回的字典里包含三个关键部分input_ids是文本在词表中的索引序列attention_mask标记哪些位置是真实文本1哪些是padding0BERT模型还会用到token_type_ids来区隔两个句子单文本分类任务用不到但接口默认会返回。我们再看怎么把原始DataFrame转成PyTorch可用的Datasetimport torch from torch.utils.data import Dataset class TextClassificationDataset(Dataset): def __init__(self, texts, labels, tokenizer, max_len128): 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 self.texts[idx] label self.labels[idx] encoded self.tokenizer.encode_plus( text, max_lengthself.max_len, paddingmax_length, truncationTrue, return_tensorspt ) return { input_ids: encoded[input_ids].squeeze(0), attention_mask: encoded[attention_mask].squeeze(0), labels: torch.tensor(label, dtypetorch.long) }注意return_tensorspt返回的是带批次维度的张量所以要squeeze(0)去掉第一维否则后续DataLoader会多叠一层维度训练时直接报形状错误。2.4 数据划分与标签编码训练集、验证集、测试集的划分比例我习惯用8:1:1。划分时一定要做分层采样保证每个类别在三个集合里的分布比例一致不然可能出现某些小类在测试集里一个样本都没有的情况。from sklearn.model_selection import train_test_split from sklearn.preprocessing import LabelEncoder label_encoder LabelEncoder() df[label_id] label_encoder.fit_transform(df[label]) train_df, temp_df train_test_split( df, test_size0.2, stratifydf[label_id], random_state42 ) valid_df, test_df train_test_split( temp_df, test_size0.5, stratifytemp_df[label_id], random_state42 )LabelEncoder会自动把字符串标签变成连续的整数id保存这个encoder预测新数据时要用同一个encoder做转换标签映射不能乱。3. 核心代码实现模型加载与训练流程3.1 加载BERT预训练模型这一行代码是整个项目的灵魂from transformers import BertForSequenceClassification model BertForSequenceClassification.from_pretrained( bert-base-chinese, num_labelsnum_classes )BertForSequenceClassification是transformers封装好的“BERT 分类头”结构预训练权重加载后分类头层是随机初始化的训练过程会更新整个模型的参数。默认情况下所有参数都会被微调这在小数据集上容易过拟合后面我会讲怎么应对。如果你的显存不够可以考虑冻结部分BERT层只训练分类头和最后几层for param in model.bert.parameters(): param.requires_grad False for param in model.bert.encoder.layer[-2:].parameters(): param.requires_grad True这个操作能节省大概30%的显存代价是效果可能会有轻微下降。个人建议数据量小几万条以内可以试数据量大的话还是全参数微调。3.2 训练参数选择与计算训练参数这块我踩过不少坑直接给一组实测比较稳的配置参数推荐值说明batch_size16或32看显存大小16更稳learning_rate2e-5BERT微调最常用的学习率区间max_len128文本长度超了会被截断不够用再调大epochs3~5BERT微调不需要太多轮次多了过拟合warmup_steps前10%步数先小学习率热身再升到设定值weight_decay0.01正则化抑制过拟合为什么BERT微调的学习率这么小因为预训练模型已经学到了很成熟的语义特征微调用大学习率会把学到的知识“冲掉”用2e-5到5e-5这个区间比较安全。我的经验是学习率首选2e-5如果训练震荡损失忽高忽低就降到1e-5收敛太慢就提到3e-5不要超过5e-5。训练循环可以这样写from torch.utils.data import DataLoader from transformers import AdamW, get_linear_schedule_with_warmup train_loader DataLoader(train_dataset, batch_size16, shuffleTrue) valid_loader DataLoader(valid_dataset, batch_size32, shuffleFalse) optimizer AdamW(model.parameters(), lr2e-5, weight_decay0.01) total_steps len(train_loader) * epochs scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_stepsint(0.1 * total_steps), num_training_stepstotal_steps ) for epoch in range(epochs): model.train() 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, attention_maskattention_mask, labelslabels) loss outputs.loss logits outputs.logits loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step() optimizer.zero_grad()你发现我加入了一行clip_grad_norm_这个操作是把梯度的范数限制在1.0以内防止梯度爆炸导致训练崩掉。BERT这种深层模型在微调初期很容易梯度波动大加这一行能显著提升训练稳定性属于小细节大作用。3.3 评估与预测验证集准确率是衡量模型效果最直观的指标但光看准确率不够多分类问题一定要看每个类别的准确率和召回率不然某类样本少时整体准确率看起来不错但小类效果差得一塌糊涂。from sklearn.metrics import classification_report, accuracy_score def evaluate(model, data_loader, device): model.eval() predictions [] true_labels [] 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[labels].to(device) outputs model(input_ids, attention_maskattention_mask) preds torch.argmax(outputs.logits, dim1) predictions.extend(preds.cpu().tolist()) true_labels.extend(labels.cpu().tolist()) acc accuracy_score(true_labels, predictions) report classification_report(true_labels, predictions, target_nameslabel_encoder.classes_) return acc, report分类报告会打印出每个类别的precision、recall、f1-score和样本数。拿到这份报告后重点看哪个类别的f1特别低低是因为类别样本太少还是因为文本语义本来就和别的类接近。这个分析直接决定你后续怎么对症下药。4. 常见问题与排查技巧实录4.1 显存不足与训练过慢BERT-base中文版参数量大约1.02亿用GPU训练是必须的纯CPU不是不能跑但一个epoch可能要跑几小时。显存不够报错信息一般是CUDA out of memory。解决办法按优先级排序把batch_size从32降到16再降到8把max_len从128降到64前提是文本真的没那么长使用梯度累积模拟更大的batch_size而不用额外占显存梯度累积代码很简单accumulation_steps 4 for step, batch in enumerate(train_loader): loss outputs.loss / accumulation_steps loss.backward() if (step 1) % accumulation_steps 0: optimizer.step() scheduler.step() optimizer.zero_grad()这样实际等效的batch_size是16乘以4等于64但显存占用还是16份的量。需要注意学习率可能需要适当调大一点来配合更大的有效batch。4.2 过拟合与验证集准确率上不去BERT在小数据集上特别容易过拟合典型表现是训练集准确率一路冲到99%验证集却卡在85%左右不动。我的经验有三个应对手段。第一增加早停机制early stopping监听验证集损失连续几个epoch不下降就停止训练并恢复最佳模型权重。第二早停前提是先把epoch数设大我通常设10让模型有足够空间去慢慢收敛。第三可以给分类头那一层加一点dropoutmodel.classifier.dropout torch.nn.Dropout(p0.3)如果你的训练数据只有几百条那模型效果大概率上不去这时候最有效的方法是找更多的数据或者用数据增强技术比如中文文本的回译增强翻译成英文再翻回来。不要指望模型在超小数据上还能有惊艳表现这不现实。4.3 预测阶段常见的坑训练完模型做预测时有两个坑要特别提醒。第一个是输入格式问题。新数据的文本必须通过同一个tokenizer、用同样的max_len处理不能直接喂原始字符串给模型否则会报维度不对或直接预测得很离谱。第二个是标签映射问题。模型输出的是softmax概率分布argmax得到的是类别id必须用之前保存的label_encoder.inverse_transform把id转回原来的字符串标签。实际预测代码def predict(model, tokenizer, text, device, max_len128): model.eval() encoded tokenizer.encode_plus( text, max_lengthmax_len, paddingmax_length, truncationTrue, return_tensorspt ) input_ids encoded[input_ids].to(device) attention_mask encoded[attention_mask].to(device) with torch.no_grad(): logits model(input_ids, attention_maskattention_mask).logits pred_id torch.argmax(logits, dim1).item() probability torch.softmax(logits, dim1).max().item() return label_encoder.inverse_transform([pred_id])[0], probability返回概率值很有用比如做内容审核时概率低于0.6的样本可以标为“待人工复核”而不是直接给出预测标签。4.4 多分类混淆矩阵的可视化调优阶段我强烈建议把混淆矩阵画出来。它能直观暴露模型容易把哪两个类搞混比如“体育”和“娱乐”经常混淆那就说明这些样本在语义上太接近需要考虑扩充特征或者换更大的模型。from sklearn.metrics import confusion_matrix import matplotlib.pyplot as plt import itertools def plot_confusion_matrix(y_true, y_pred, classes, titleConfusion Matrix): cm confusion_matrix(y_true, y_pred) plt.figure(figsize(10, 8)) plt.imshow(cm, interpolationnearest, cmapplt.cm.Blues) plt.title(title) plt.colorbar() tick_marks range(len(classes)) plt.xticks(tick_marks, classes, rotation45) plt.yticks(tick_marks, classes) thresh cm.max() / 2. for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])): plt.text(j, i, format(cm[i, j], d), horizontalalignmentcenter, colorwhite if cm[i, j] thresh else black) plt.ylabel(真实标签) plt.xlabel(预测标签) plt.tight_layout() plt.savefig(confusion_matrix.png, dpi200)可视化做得越清楚你越容易定位问题。有一次我排查一个新闻分类模型看到“财经”和“金融”类之间混淆特别严重后来发现是因为训练数据里这两个类别的定义本身就重合了重新清洗标注之后模型效果明显提升。数据质量永远是模型效果的上限。5. 模型优化方向与扩展思路5.1 换更强的中文预训练模型如果bert-base-chinese的效果不满足需求transformers库里有几个可以直接替换的升级版hfl/chinese-roberta-wwm-ext、hfl/chinese-bert-wwm-ext、uer/roberta-base-finetuned-jd-binary-chinese。它们用了全词掩码或更大规模的语料训练很多任务上比原始BERT有提升。替换方法就是改模型名model BertForSequenceClassification.from_pretrained( hfl/chinese-roberta-wwm-ext, num_labelsnum_classes ) tokenizer BertTokenizer.from_pretrained(hfl/chinese-roberta-wwm-ext)注意RoBERTa系列没有token_type_ids但BertForSequenceClassification内部会处理你不需要改代码。这类模型下载时会自动缓存到~/.cache/huggingface目录如果公司内网不能直连HuggingFace可以手动下载模型文件后从本地加载model BertForSequenceClassification.from_pretrained(/your/local/path/model, num_labelsnum_classes)5.2 处理类别不平衡问题类别不平衡在真实业务场景里非常常见金融投诉文本里“理财纠纷”类占了70%剩下30%分散在十几类里。直接训练会得到“看似准确率高、实则小类全废”的模型。两个方向可以同时用。损失函数层面在BERT模型输出logits后不直接使用默认的cross entropy换成带权重的torch.nn.CrossEntropyLoss(weightclass_weights)每个类别的权重可以设为样本数的倒数。代码上需要把模型的loss忽略自己算loss来反向传播class_weights torch.tensor([...], devicedevice) # 按样本数倒数归一化 loss_fn torch.nn.CrossEntropyLoss(weightclass_weights) outputs model(input_ids, attention_maskattention_mask) loss loss_fn(outputs.logits, labels)数据层面对少数类做上采样简单重复复制或者用过采样方法如SMOTE在向量空间内插值生成新样本。NLP数据上SMOTE效果因数据而异需要实验验证。我对类别数少且样本也少的任务优先用带权重的损失函数改动最小且稳定。5.3 从研究到工程部署的思考训练好的模型最终要部署上线就涉及模型导出和推理加速。PyTorch环境下常规方案是将模型转成ONNX或者TorchScript。ONNX导出后可以配合ONNX Runtime推理速度大约能提升1.5到3倍具体取决于硬件优化。import torch from transformers import BertForSequenceClassification model BertForSequenceClassification.from_pretrained(your_model_dir) model.eval() dummy_input { input_ids: torch.randint(0, 100, (1, 128)), attention_mask: torch.ones(1, 128, dtypetorch.long) } torch.onnx.export( model, (dummy_input[input_ids], dummy_input[attention_mask]), model.onnx, input_names[input_ids, attention_mask], output_names[logits], dynamic_axes{input_ids: {0: batch_size}, attention_mask: {0: batch_size}, logits: {0: batch_size}}, opset_version12 )dynamic_axes这个参数必须设置不然导出的模型batch_size固定为1线上请求并发时会很吃亏。部署时还可以用TensorRT做进一步加速如果团队有精力折腾的话。6. 踩坑总结与个人心得做这个项目的过程中有几个坑印象特别深。第一个是早期我把学习率设成了1e-3想着优化器是AdamW应该能自适应结果训练损失疯狂震荡完全收敛不下去。后来才意识到预训练模型的参数美丽之处已经在预训练阶段“长好了”微调只适合缓慢调整2e-5这个量级才是正确区间。第二个坑是分词器版本不一致。有一次我换了transformers版本没有重新下载最新分词器结果老模型加载新分词器tokenizer.encode(科技)出来的id和训练时完全不一样预测效果直接掉到随机水平。这个排查花了两个多小时最后重新下载对应版本的分词器才解决。后来我学乖了训练完一定会把tokenizer和模型一起保存到目录model.save_pretrained(./saved_model) tokenizer.save_pretrained(./saved_model)再加载时直接用from_pretrained(./saved_model)保证分词器和模型版本严格匹配从那以后再没出过这个问题。第三个心得和数据处理相关。我实测下来文本清理环节对中文BERT的影响没有想象中那么大BERT对中文的字级泛化能力很强轻微的标点缺失和错别字影响有限。但这不代表可以不清理数据。关键是控制文本长度max_len从64加到128能带来稳定的效果提升但再往上加到256提升就非常有限了而训练时间几乎翻倍。128是一个性价比很高的值。最后再分享一个调优技巧训练完成后不要直接拿测试集做最后评估先反复看验证集的分类报告。把验证集里预测错的样本打印出来一条一条看归类错误的原因。你会发现有些是标注错误数据问题有些是语义确实模糊模型能力上限还有些是某种句式模型没见过数据多样性不足。分析完这个你的下一步优化方向就会非常清晰而不是盲目调参。BERT中文文本分类这件事技术上不复杂但每个环节都有值得打磨的细节。把数据、参数、评测这三块做扎实了模型效果自然就上来了。本文还有配套的精品资源点击获取