简介基于pytorch-transformers实现的BERT中文文本分类项目定位于帮助NLP开发者与学生快速掌握预训练模型在中文场景下的微调流程解决中文新闻标题多分类任务。项目提供完整可运行代码覆盖数据预处理、模型训练、验证与测试环节并附带从THUCNews抽取的20万条新闻标题数据涵盖财经、房产、股票等10个类别按18万/1万/1万划分为训练集、验证集和测试集可直接体验BERT对中文文本的分类效果。压缩包共29个文件以Python源码、txt数据文件、xml配置、bin模型权重和png结果图为主整体约732.57MB目录结构清晰便于对照README.docx理解各模块用途。已有3699人学习该资源适合作为中文NLP分类任务的基础实现参考。 最近在处理一个中文文本分类需求时我把基于 pytorch-transformers 实现的 BERT 中文文本分类代码整套流程重新捋了一遍。从数据准备、模型微调到最后的评估与部署推理每一步都实测跑通也踩了不少坑。如果你正打算用 BERT 做中文分类任务或者看了不少零散教程还不知道代码怎么串起来这篇博文可以直接帮你省出一周的摸索时间。所谓的 pytorch-transformers其实就是 HuggingFace 的 transformers 库早期叫法现在新项目直接装 transformers 就行。无论叫哪个名字核心思路不变加载一个已经在海量中文语料上预训练好的 BERT 模型在少量标注数据上做微调然后用于你手头的分类任务。我下面会把这套代码拆开来讲包括环境版本坑、数据预处理、Dataset 与 DataLoader 的构建、训练循环里的参数选择以及最后评估和模型保存所有代码都是可以直接照抄改用的水平。1. 项目整体设计与思路拆解1.1 为什么用 BERT 而不是传统文本分类方案做中文文本分类可选的技术路线很多。最简单的是 TF-IDF 加逻辑回归或朴素贝叶斯稍微进阶一点的用 Word2Vec 词向量加 TextCNN 或 BiLSTM。但在实际项目里这些方案我从准确率、泛化能力、落地成本三个维度对比过BERT 微调在绝大部分中文本分类场景下都是性价比最高的基线方案。传统 TF-IDF 方案最大的问题是不理解语义。它把文本当成一个词袋只统计词频和逆文档频率同义词、反讽、语序颠倒全都处理不了。比如“这个电影一点都不难看”和“这个电影很难看”字面差异只有“不”一个词但情感方向完全相反。TF-IDF 通常会把这两句话的特征向量算得很接近分类器自然容易分错。TextCNN 和 BiLSTM 这类深度学习模型比词袋强不少能抓到局部 n-gram 特征和上下文信息但它们的词向量本身是静态的同一个词在不同语境下始终是同一个向量无法解决一词多义问题。而 BERT 的核心优势在于它通过 Transformer 的双向编码结构在预训练阶段就学会了给每个词根据上下文赋予动态向量表示。同样是“苹果”在“苹果很好吃”和“苹果发布了新手机”里向量完全不同。这个特性对中文分类任务尤其重要因为中文的表达更加依赖语境。另外一个很实际的因素是工程成本。BERT 微调不需要你从零训练语言模型所有语言知识已经包含在预训练权重里。你只需要准备几千到几万条标注数据在单张 GPU 上花几十分钟到几个小时就能得到一个效果还不错的分类模型。相比自己从头训练一个深层网络这个成本低得多也稳定得多。1.2 项目流程与代码结构设计整个项目可以拆成五个阶段数据准备、文本编码、模型加载、训练微调、评估推理。对应到代码上我用一个很简洁的工程结构来组织bert_text_classification/ ├── data/ │ ├── train.csv # 训练集text,label两列 │ ├── dev.csv # 验证集 │ └── test.csv # 测试集 ├── main.py # 训练主脚本 ├── predict.py # 推理脚本 ├── requirements.txt # 依赖文件 └── README.md这个结构不复杂但足够应付中小规模的中文分类项目。数据统一用 CSV 格式存储每行是一条文本加一个整数标签简单直观Pandas 直接读入。训练脚本负责加载数据、构建 Dataset、初始化模型、跑训练循环并在每个 epoch 结束后在验证集上评估一次把效果最好的模型保存下来。推理脚本则独立运行加载保存好的模型对单条文本做预测方便后续接到 Flask 或 FastAPI 服务里。设计上有一个关键取舍不把训练和推理耦合在同一个脚本里。很多新手喜欢把训练完的模型直接在当前进程里做预测这样看似方便但实际部署时会很别扭因为生产环境通常只需要推理不需要再加载庞大的训练数据。拆成两个脚本训练和推理的环境可以分离模型文件也能独立管理。2. 环境准备与核心依赖踩坑2.1 从 pytorch-transformers 到 transformers 的版本认知我之所以强调 pytorch-transformers 这个老名字是因为很多老教程、老代码片段里还在用它新手照着抄容易出现包名不存在的报错。更准确的版本沿革是这样的HuggingFace 最早确实发布过单独支持 PyTorch 的 pytorch-transformers 包后来又出了单支持 TensorFlow 的 pytorch-transformers 的姊妹包再后来统一合并成了 transformers 这个仓库同时支持 PyTorch、TensorFlow 和 JAX。旧的 pip 安装命令是pip install pytorch-transformers导入时写from pytorch_transformers import BertTokenizer带下划线。现在的新版本安装命令是pip install transformers导入时写from transformers import BertTokenizer没有下划线。如果你从旧帖子复制代码记得把pytorch_transformers全局替换成transformers。版本兼容上我遇到过最坑的一件事是 transformers 4.x 之后 API 有一些不兼容改动。比如旧版本里model.bert可以直接访问底层 BERT 编码器新版本里如果你用BertForSequenceClassification加载模型访问方式依然是model.bert但如果你用更通用的AutoModelForSequenceClassification底层模型名可能是model.bert也可能根据配置变化直接硬编码访问就会出错。建议代码里不要依赖这类内部属性如果想拿到句向量用模型输出的 raw hidden states 才是最稳的方式。2.2 数据格式与预处理细节数据样例我这里给一个三分类的片段实际项目里二分类、多分类都可以套用text,label 这台手机性价比很高屏幕显示效果一流。,0 快递等了十天都没到客服态度还很差。,1 剧情平淡演员演技在线但整体节奏拖沓。,2第一列文本第二列整数标签。类别标签最好从 0 开始连续编号因为模型的num_labels参数直接对应你传入的类别数。如果数据集里的标签是中文比如“正向”“负向”需要先做一次映射把中文转成整数。读入数据后我强烈建议先做两件事查看文本长度分布和检查标签分布。长度分布决定max_length参数怎么设。用 BERT 做分类文本会先被 tokenizer 切成 token中文 BERT 的 vocab 是按字级别构建的切出来的 token 数量大约等于汉字字符数加两个特殊 token。如果你发现 95% 以上的文本都在 100 个字以内max_length128就够了如果文本普遍很长再考虑 256 或 512。检查标签分布主要是为了发现类别不平衡问题。如果某个类别只占 5% 的数据模型很容易把所有样本都预测成大类。最简单的处理方式是在损失函数里给每个类别加权重PyTorch 的CrossEntropyLoss本身支持weight参数按类别样本数的倒数归一化即可。我后面在训练代码里也会给出具体用法。预处理这一步有一个新手常犯的错误就是对中文文本做分词。BERT 的中文 tokenizer 不需要 jieba 分词它会自动按字符切分。你如果先用 jieba 把文本切成词再用空格拼起来喂给 BERT反而会引入错误的切分边界因为 BERT 的词表是基于字的你强行切成词会让它把不认识的词拆成奇怪的 token 组合语义信息反而受损。我实测下来直接输入原始文本去掉多余换行和特殊符号效果是最好的。3. 核心代码实现与参数选择3.1 加载 tokenizer 与预训练模型整个项目最关键的一步就是把预训练模型和 tokenizer 从 HuggingFace 模型库拉到本地。代码如下from transformers import BertTokenizer, BertForSequenceClassification model_name bert-base-chinese tokenizer BertTokenizer.from_pretrained(model_name) model BertForSequenceClassification.from_pretrained( model_name, num_labels3, )这里的bert-base-chinese是 Google 发布的 BERT 中文基础版本也是目前最常用的中文预训练模型词表大小约 21128编码器 12 层隐藏层维度 768。如果你的数据集规模不大用这个就够了。HuggingFace 上还有哈工大讯飞联合发布的chinese-bert-wwm-ext和chinese-roberta-wwm-ext这两个模型用了全词掩码策略在部分中文任务上表现会略好一些但需要从国内镜像下载网络好的话也可以直接尝试。第一次运行上面的代码会自动从模型库下载权重文件大概 400MB 左右等待时间取决于网络。建议第一次先把模型跑通确认没有任何报错再进入正式训练。如果你的网络下载困难可以先用国内镜像地址比如设置环境变量HF_ENDPOINThttps://hf-mirror.com再运行代码。加载模型后建议立刻打印一下模型参数量total_params sum(p.numel() for p in model.parameters()) trainable_params sum(p.numel() for p in model.parameters() if p.requires_grad) print(fTotal params: {total_params / 1e6:.2f}M) print(fTrainable params: {trainable_params / 1e6:.2f}M)bert-base-chinese 的参数量约 102M微调时所有参数都会参与梯度更新。如果你机器显存很紧张后面我会讲冻结部分层只训分类头的省显存方案但一般情况下全量微调效果最好。3.2 Dataset 与 DataLoader 构建PyTorch 训练数据需要包装成 Dataset 和 DataLoader。这一步的核心逻辑是把原始文本通过 tokenizer 编码成模型需要的三个输入input_ids、attention_mask、labels然后按批次抛给模型。from torch.utils.data import Dataset import torch class BertTextDataset(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 str(self.texts[idx]) label self.labels[idx] inputs self.tokenizer( text, max_lengthself.max_len, paddingmax_length, truncationTrue, return_tensorspt, ) return { input_ids: inputs[input_ids].flatten(), attention_mask: inputs[attention_mask].flatten(), labels: torch.tensor(label, dtypetorch.long), }这里有几个细节要重点讲。paddingmax_length表示所有样本都补齐到max_len长度。为什么不建议用paddinglongest因为那样每个 batch 里 tokenizer 会按当前批次内最长文本补齐不同 batch 之间长度不同虽然训练时能省一点显存但数据处理会稍微慢一些而且推理时也要跟着动态算。max_length128固定长度在工程上最简单batch 的张量尺寸始终一致。truncationTrue表示文本超长时直接截断。BERT 本身最大支持 512 个 token超过部分会被丢弃。对大多数短文本分类任务来说128 已经够用。我实测过电商评论数据max_length128和max_length512的准确率差距不到 0.5 个百分点但训练速度差了一倍以上所以短文本分类优先选 128。attention_mask是很多新手容易忽略的部分。它的作用是指示模型哪些位置是真实的文本 token哪些位置是为了对齐长度补上去的 padding token。模型计算注意力分数时padding 位置会通过 mask 置为负无穷相当于告诉注意力机制不关注这些位置。如果你手动构造 input tensor 但不传 attention_mask模型默认全为 1把 padding token 也当成真实文本参与注意力计算训练效果会明显下降。DataLoader 部分from torch.utils.data import DataLoader train_loader DataLoader( train_dataset, batch_size16, shuffleTrue, num_workers2, ) val_loader DataLoader( val_dataset, batch_size64, shuffleFalse, num_workers2, )训练集shuffleTrue打乱数据顺序避免模型学到样本顺序带来的虚假模式。验证集不需要打乱。num_workers是数据加载的并行进程数Linux 下设置成 4 或 8 效果更好Windows 下设置成 0 或 2否则会有子进程报错。3.3 训练循环与学习率调度训练循环是核心中的核心。BERT 微调和从头训练一个网络在参数配置上有非常大的区别最关键的就是学习率。from transformers import AdamW, get_linear_schedule_with_warmup epochs 5 total_steps len(train_loader) * epochs optimizer AdamW(model.parameters(), lr2e-5, weight_decay0.01) scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_stepsint(total_steps * 0.1), num_training_stepstotal_steps, )BERT 微调学习率通常设置在1e-5到5e-5之间比你在普通网络上用的1e-3要小一个数量级。原因很简单预训练模型已经在一个非常大的语料上学到了成熟的语言表示你的微调数据只是在这个基础上做小幅修正。学习率过大一步就可能把预训练学到的表征冲掉出现训练几天 loss 不降反升的诡异现象。我最初做的时候吃过这个亏学习率设成1e-4模型在验证集上的表现还不如随机猜测。后来把学习率降到2e-5效果立刻就正常了。warmup机制也很重要。它的作用是前 10% 的步数里把学习率从 0 线性增加到目标值让模型在训练初期不会因为参数剧烈震荡而破坏预训练表示的稳定性。warmup 之后学习率再线性衰减到 0保证后期收敛稳定。这个策略已经是预训练模型微调的事实标准不建议改动。训练循环主体from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) best_acc 0.0 for epoch in range(epochs): model.train() total_loss 0.0 for batch in tqdm(train_loader, descfEpoch {epoch 1}/{epochs}): input_ids batch[input_ids].to(device) attention_mask batch[attention_mask].to(device) labels batch[labels].to(device) outputs model( input_idsinput_ids, attention_maskattention_mask, labelslabels, ) loss outputs.loss total_loss loss.item() loss.backward() optimizer.step() scheduler.step() optimizer.zero_grad() avg_loss total_loss / len(train_loader) print(fEpoch {epoch 1} - train loss: {avg_loss:.4f}) val_acc evaluate(model, val_loader, device) print(fValidation accuracy: {val_acc:.4f}) if val_acc best_acc: best_acc val_acc model.save_pretrained(best_model) tokenizer.save_pretrained(best_model)每个 epoch 结束后在验证集上跑一次评估记录最佳准确率并保存对应的模型权重这样即便后几个 epoch 过拟合了你也还能拿回表现最好的模型。评估函数我单独写在下面。损失层面如果你的数据类别不平衡建议在模型输出上用CrossEntropyLoss重新计算from torch.nn import CrossEntropyLoss loss_fct CrossEntropyLoss(weightclass_weights.to(device)) outputs model(input_idsinput_ids, attention_maskattention_mask) logits outputs.logits loss loss_fct(logits.view(-1, num_labels), labels.view(-1))class_weights按 1 除以每个类别的样本数再归一化即可。实际项目中我遇到过某个类别只占 8% 的数据不加权重时这个类别几乎全被预测错加了权重后 F1 直接提升 15 个百分点以上效果非常显著。4. 评估、保存与推理4.1 评估指标与混淆矩阵文本分类任务不能只看准确率尤其是类别分布不均匀的时候准确率会极具欺骗性。比如 90% 的样本属于 A 类你全预测 A 类也能拿 90% 准确率但这个模型毫无意义。正确的做法是同时看精准率、召回率和 F1。from sklearn.metrics import classification_report, confusion_matrix import numpy as np def evaluate(model, dataloader, device): model.eval() predictions [] true_labels [] with torch.no_grad(): for batch in dataloader: input_ids batch[input_ids].to(device) attention_mask batch[attention_mask].to(device) labels batch[labels].to(device) outputs model(input_idsinput_ids, attention_maskattention_mask) logits outputs.logits preds torch.argmax(logits, dim-1) predictions.extend(preds.cpu().numpy()) true_labels.extend(labels.cpu().numpy()) print(classification_report(true_labels, predictions, digits4)) print(Confusion Matrix:) print(confusion_matrix(true_labels, predictions)) return np.mean(np.array(predictions) np.array(true_labels))classification_report会输出每个类别的精准率、召回率和 F1以及宏平均和加权平均方便你判断模型对少数类别的处理能力。混淆矩阵能直观看出哪些类别之间容易混淆比如“正向”和“中向”之间经常被搞混那就说明两者的语义边界确实不够清晰可能需要在数据层面补充更多差异性的标注样本。这里千万不要省略model.eval()和torch.no_grad()。eval()会关闭 Dropout 等训练时才生效的层no_grad()关闭梯度计算。如果你漏掉这两行推理时显存占用会高很多而且输出结果也不对因为 Dropout 还在随机丢弃神经元第二次跑同一个样本结果都不一样。4.2 模型保存与在线推理训练完成后模型和 tokenizer 各自保存到目录。注意 tokenizer 也必须一起保存因为新版本的 tokenizer 会保存特殊 token 映射如果你只保存模型权重而丢了 tokenizer推理时你还得从远端的bert-base-chinese重新下载 tokenizer走一遍完全一样的初始化流程麻烦不说还可能引入版本不一致的问题。model.save_pretrained(best_model) tokenizer.save_pretrained(best_model)推理脚本predict.py的核心逻辑如下from transformers import BertTokenizer, BertForSequenceClassification import torch model_path best_model tokenizer BertTokenizer.from_pretrained(model_path) model BertForSequenceClassification.from_pretrained(model_path) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) model.eval() def predict(text): inputs tokenizer( text, max_length128, paddingmax_length, truncationTrue, return_tensorspt, ) input_ids inputs[input_ids].to(device) attention_mask inputs[attention_mask].to(device) with torch.no_grad(): logits model(input_idsinput_ids, attention_maskattention_mask).logits probs torch.softmax(logits, dim-1) pred_idx torch.argmax(probs, dim-1).item() return pred_idx, probs[0][pred_idx].item() if __name__ __main__: test_text 这个产品质量太差了用了两次就坏了。 label, conf predict(test_text) print(f预测标签: {label}, 置信度: {conf:.4f})推理时关闭梯度用softmax把 logits 转成概率分布这样既能拿到预测标签也能拿到置信度方便后续做阈值过滤把低置信度的预测交给人工处理。5. 常见问题与排查技巧实录5.1 训练不收敛与损失异常我最常被问到的就是训练时 loss 不下降或者直接变 NaN。这类问题排查顺序很有规律这里整理成表格方便对照现象可能原因解决方案Loss 一直很大不下降学习率过大或过小调到 2e-5 到 5e-5 区间Loss 变 NaN梯度爆炸加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)训练 loss 降了验证 loss 反弹过拟合增加数据量、早停、减小模型规模所有预测都是同一个标签类别严重不平衡给损失函数加权重或改用 Focal Loss验证集表现极差但训练集正常数据划分有泄漏检查是否存在同一文本同时出现在训练和验证集梯度裁剪是我强烈建议直接加上的一行代码。BERT 的深层结构在训练初期很容易出现梯度范数过大导致的参数震荡加上clip_grad_norm_就像给训练过程装了个保险丝几乎没有任何副作用。5.2 显存不足与训练加速如果你只有一块消费级显卡比如 8GB 显存跑 bert-base-chinese 最大 batch_size 16文本长度 128 压力不大但如果文本长度设成 512显存占用会直线上升很容易 OOM。遇到显存不足不要直接换更大的显卡有几个更经济的方案第一降低 batch_size 并用梯度累积补偿。BERT 在小 batch 下训练不稳定梯度累积可以在不增加显存的情况下模拟大 batch 的效果。把整个过程拆成若干微批次每个批次反向传播计算梯度但先不更新参数累积若干批后再统一优化一次。accumulation_steps 4 optimizer.zero_grad() for step, batch in enumerate(train_loader): outputs model(**batch) loss outputs.loss / accumulation_steps loss.backward() if (step 1) % accumulation_steps 0: optimizer.step() scheduler.step() optimizer.zero_grad()第二开启混合精度训练。PyTorch 的自动混合精度可以在几乎不影响精度的情况下把显存占用降低 40% 以上。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): outputs model(input_idsinput_ids, attention_maskattention_mask, labelslabels) loss outputs.loss scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()第三冻结 BERT 底层网络层只训练分类头和最后几层 Transformer 层。这个方法在标注数据较少时还能降低过拟合风险具体做法是把model.bert.encoder.layer的前八层设成requires_gradFalse显存占用能再降 20% 左右。5.3 长文本截断与信息保留BERT 的max_position_embeddings是 512超过 512 个 token 的文本物理上就塞不进去。这是模型结构决定的硬限制不是参数调优能解决的。短文本分类基本不用担心但如果你处理的是新闻、长评论、法律文书这类长文本截断策略直接影响效果。最简单的方案是直接截断前 512 个 token。但遇到重要信息在文本后半段的情况就麻烦了比如一个用户先写了一大堆铺垫最后一句才是核心观点你只保留前半段就把核心信息丢掉了。实际项目中我试过两种改良方案第一种是头尾保留策略取前 400 个 token再取后 112 个 token中间丢弃。这样既保留了开头通常是主题句又保留了结尾通常是总结句信息覆盖率大幅提升。用 tokenizer 做不到直接这样做需要在编码前手动切分文本。第二种是分段 聚合策略把长文本按滑动窗口切成长度不超过 500 的多个片段每个片段单独过模型得到句向量然后对所有片段向量做平均池化或者最大池化再送进分类器。这种方法能利用全文信息但训练和推理成本会成倍增加一般在超长文档分类场景才值得上。我自己最常采用的是头尾保留策略代码改动不到十行效果提升非常明显。以一份 800 字的长评为例直接截断前 512 个 token 准确率 81%头尾保留后提升到 86%这个差距在业务应用里已经很大了。整套代码我目前已经用在了电商评论情感分类、舆情正负面判断、工单自动打标等多个实际项目中单条推理耗时在 GPU 上大约是 10 毫秒级别CPU 上约 100 毫秒完全能支撑绝大多数线上服务。如果你要做的是实时性要求极高的场景还可以再把模型导出成 ONNX 格式推理延迟还能再压一半。文本分类这条路BERT 微调是起点但绝对够用很久。本文还有配套的精品资源点击获取
