
简介面向自然语言处理学习者的 BERT 中文文本分类实战工程基于 pytorch-transformers 构建解决中文新闻标题多分类问题。数据采用 THUCNews 20万条新闻标题覆盖财经、房产等10个类别按18万/1万/1万切分训练、验证与测试集并含数据预处理、模型训练与推理验证等完整环节。资源共29个文件以 Python 脚本、预训练 BERT 模型文件bin/json/txt、结果图片与说明文档为主压缩包约732.57MB。目录提供 README、运行结果记录及可视化图片便于复现实验与对照改进。已有3700人浏览学习。下载即可获得数据处理脚本、模型训练入口、配置参数、预训练权重与评估结果适合基于真实中文数据集快速上手 BERT 分类的课程设计或科研场景。1. BERT 中文文本分类一份能直接改业务场景的 PyTorch 实现手里压着一堆用户反馈文本想按售后、物流、产品、投诉几类自动打标这是很多人第一次动“用 BERT 做中文文本分类”念头的场景。我当初接手类似需求时从装环境到跑通第一版只花了不到两天用的就是这份基于 pytorch-transformers如今叫 transformers 生态封装的中文文本分类代码。它不是算法玩具而是一套完整可落地的工程流程数据预处理、模型加载、训练、评估、预测都拆成了清晰模块预训练权重直接用 bert-base-chinese你只需要准备一份标准格式的文本数据文件就能训练出属于自己的分类模型。适合刚接触 NLP 方向、想把 BERT 用在真实数据上的开发者也适合手头有标注数据、想快速验证分类效果的工程团队——这份代码的价值在于把 BERT 从“黑匣子”变成了“可改的流水线”。2. 技术选型与整体架构从 pytorch-transformers 到一条可替换的训练主线2.1 版本选择与依赖环境训练 BERT 中文分类模型PyTorch 是绝对主流因为 Hugging Face 的 transformers 库对 PyTorch 的兼容性最完整调试和断点续训都方便。标题里提到的 pytorch-transformers 其实就是 transformers 的前身早期版本叫这个名字后来库合并统一成了 transformersAPI 基本兼容。所以你现在装依赖直接装新库就行代码里导包也是from transformers import AutoTokenizer, BertForSequenceClassification这一类写法。我一般建议用 Python 3.8 以上的环境PyTorch 装 2.x 版本即可transformers 装 4.x 版本。如果跑的是 GPU 环境注意 CUDA 版本要和 PyTorch 对应这个细节决定了你torch.cuda.is_available()是不是返回 True。安装依赖通常一条命令搞定pip install torch transformers datasets accelerate scikit-learn pandas其中datasets不是必须的但用它加载 CSV 和做 train/valid 切分很方便accelerate会在你后续想改成多卡训练时省很多事。安装完之后建议先跑一行验证from transformers import AutoTokenizer tok AutoTokenizer.from_pretrained(bert-base-chinese) print(tok(你好世界))能正常打印 token 序列说明环境没问题。第一次运行会自动下载 bert-base-chinese 的模型文件到缓存目录这个过程取决于网络条件模型文件大概几百 MB建议提前单独下载保存到本地目录后面加载时直接用本地路径能省不少时间也避免每次启动都去检查远端版本。2.2 目录结构与职责划分这份代码的目录结构是典型的“训练/推理分离”布局每个文件责任单一。我拆给你看bert_chinese_classification/ ├── data/ # 数据集目录放 train.csv / dev.csv / test.csv ├── model_save/ # 训练好的模型和 tokenizer 保存位置 ├── config.py # 全局配置路径、超参数、标签数量 ├── data_processor.py # 数据加载、清洗、tokenize、Dataset 封装 ├── train.py # 训练主脚本 ├── predict.py # 模型预测脚本 ├── utils.py # 通用工具函数指标计算、seed 设置、模型保存 └── requirements.txt # 依赖清单为什么要拆成config.py而不是把所有参数写在train.py里因为跑实验的时候调 batch size、学习率、max length 的频率非常高如果每次都要去几千行的训练脚本里翻参数效率很低也比较容易改错。把配置独立出来之后换数据集、改标签体系都只动一个文件。数据文件格式建议统一为 CSV至少包含两列text和label。label 可以直接是中文类别名称比如“售后”“物流”代码里会自己建立映射表。训练时会自动把 label 转成 0、1、2 这样的数字索引这个映射表会保存下来预测阶段加载模型时要用到同一张表否则预测结果会错位。2.3 快速启动训练一条命令跑通数据准备好之后启动训练其实就是一行命令python train.py --do_train --epochs 3 --batch_size 16 --max_length 128跑完会在model_save/目录下生成模型权重pytorch_model.bin、配置文件config.json、词表vocab.txt还有一个记录 label 映射关系的label2id.json。接着用同一份代码做预测python predict.py --text 这个包裹三天了还没送到客服也不回消息输出大概是{label: 物流, confidence: 0.93}这样的结果。到这里一条“原始文本文件 - 训练 - 模型 - 预测”的主线就完整了。接下来我们进到每个环节内部看数据是怎么被喂给模型的。3. 数据预处理是真正的战场分词、截断与注意力掩码3.1 中文 tokenizer 的使用细节BERT 中文分类的数据处理第一个关键点是分词。很多人第一次接触 tokenizer 时会下意识想“中文是不是要先结巴分词”事实上 bert-base-chinese 的词表就是按字character-level构建的直接按 token 切分不需要额外接分词工具。你只需要把原始文本交给 tokenizer它会自动做好字粒度切分、添加[CLS]和[SEP]标记。一个特别值得注意的细节是tokenizer 返回的是三个张量input_ids、token_type_ids和attention_mask。input_ids是每个 token 在词表中的编号attention_mask标记哪些位置是真实内容值为 1、哪些位置是 padding值为 0token_type_ids在单文本分类里全为 0。我第一次写 DataLoader 时只用input_ids喂给模型结果训练 loss 一直降不下去后来才发现是 attention_mask 没传模型把 padding 部分也当成了真实语义。from transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(bert-base-chinese) def encode_text(text, max_length128): encoded tokenizer( text, max_lengthmax_length, paddingmax_length, truncationTrue, return_tensorspt ) return encoded[input_ids][0], encoded[attention_mask][0] # 示例 input_ids, attention_mask encode_text(这款产品的质量超出预期非常满意) print(input_ids[:10]) print(attention_mask[:10])max_length128表示最长保留 128 个 token超出部分截掉不足部分用 0 填充到统一长度这样 batch 内每个样本维度一致才能拼成一个矩阵送进 GPU。truncationTrue表示超长截断默认是保留开头丢弃尾部这正是很多分类效果不佳的隐患——对于长文本像“客服不回复一直到第无数次投诉”这种关键信息如果出现在结尾默认截断就把它丢掉模型可能完全看不到决定性的内容。后面避坑章节我会专门展开这个点。3.2 Dataset 类与动态 padding 策略在 PyTorch 训练链路中把数据组织成 Dataset 后交给 DataLoader是最常规的做法。关键点在 collate_fn也就是如何把一个 batch 的样本拼起来。常见做法有两种固定长度 padding 和动态 batch 内 padding。这份代码用的是固定长度 padding因为简单可控所有样本统一到max_length无论 GPU 推理还是转 ONNX 都方便。import torch from torch.utils.data import Dataset, DataLoader class BertClassificationDataset(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 str(self.texts[idx]) label self.labels[idx] encoded self.tokenizer( text, max_lengthself.max_length, 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) } def collate_fn(batch): input_ids torch.stack([item[input_ids] for item in batch]) attention_mask torch.stack([item[attention_mask] for item in batch]) labels torch.stack([item[labels] for item in batch]) return { input_ids: input_ids, attention_mask: attention_mask, labels: labels }collate_fn在这里做了三件事把 input_ids、attention_mask、labels 分别从 list of tensor 堆叠成矩阵。dtypetorch.long是必须的分类标签要用长整型否则后面计算交叉熵损失会报类型不匹配的错误。DataLoader 里设置shuffleTrue可以打乱训练数据避免模型记住样本顺序。3.3 数据清洗全半角与空白字符处理中文数据天然带着一堆脏东西全角逗号、全角括号、空格、换行、特殊符号、繁体字。如果不做清洗tokenizer 会把全角符号切成罕见 token严重干扰模型学习。常见做法是写一个清洗函数在进入 tokenizer 之前先把文本归一到标准形态。import re def clean_text(text): # 全角转半角统一逗号、括号、空格等符号 text text.replace(, ,).replace(。, .) text text.replace(, ().replace(, )) text text.replace(, :).replace(, ;) text text.replace(”, ).replace(“, ) # 去掉多余空白 text re.sub(r\s, , text).strip() # 去除特殊 URL text re.sub(rhttps?://\S|www\.\S, , text) return text这个清洗函数看起来简单却直接影响训练效果。我曾经拿一份用户评论数据不洗就训练验证集 F1 卡在 0.82后来把全角转半角加上同样的参数几分钟就涨到 0.87。原因在于全角逗号在 bert-base-chinese 的词表里对应一个很偏的 token模型要为它额外学习一套模式白白浪费容量。注意清洗动作不要在每次训练前做一次就完了要保证训练集和预测输入走的是同一条清洗逻辑否则训练时模型没见过“带全角逗号的新文本”部署之后效果会打折扣。4. 核心代码拆解从模型加载到训练循环的完整链路4.1 加载预训练模型与参数配置BERT 分类模型的核心就是BertForSequenceClassification它等于“BERT 编码器 Pooler 线性分类头”。加载时要注意两个参数num_labels决定分类头输出的类别数id2label和label2id让模型自带的分类头在推理时能输出可读的类别名称。from transformers import BertForSequenceClassification, BertConfig model BertForSequenceClassification.from_pretrained( bert-base-chinese, num_labelslen(label2id), label2idlabel2id, id2label{v: k for k, v in label2id.items()} )如果是从头微调首次加载时会看到一段提示说随机初始化了一些权重这是正常的BertForSequenceClassification的分类头是随机初始化的需要靠你的训练数据来学习。label2id必须是稳定的映射表不能每次训练都用 set 去重生成否则类别顺序乱了保存的模型下次加载时就会张冠李戴。4.2 优化器与学习率调度warmup 是必需品训练 BERT 微调任务优化器通常选择 AdamW搭配线性 warmup 调度器这也是 transformers 库的标准做法。直接上完整代码from transformers import AdamW, get_linear_schedule_with_warmup def build_optimizer(model, num_training_steps): no_decay [bias, LayerNorm.weight] optimizer_grouped_parameters [ { params: [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], weight_decay: 0.01 }, { params: [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], weight_decay: 0.0 } ] optimizer AdamW(optimizer_grouped_parameters, lr2e-5, eps1e-8) warmup_steps int(num_training_steps * 0.1) scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_stepswarmup_steps, num_training_stepsnum_training_steps ) return optimizer, scheduler为什么 bias 和 LayerNorm 权重不设 weight decay这是训练 NLP 模型的通用经验这两类参数本身是归一化相关的偏置项正则化它们容易导致训练不稳定。学习率设置 2e-5 是微调 BERT 的常用起点如果你发现训练集 loss 降得很慢可以先尝试 3e-5但一般不要超过 5e-5再大就容易出现灾难性遗忘。warmup_steps 是总步数的 10%意思是前 10% 的步数里学习率从 0 线性升到目标值这是为了防止训练初期模型参数剧烈震荡。4.3 训练主循环梯度累积与早停检查训练循环的骨架是标准的 PyTorch 流程但有几个细节值得单独说明。梯度累积是一种宽慰显存焦虑的手段如果你的 GPU 装不下 batch_size32可以设置 batch_size8梯度累积 4 步再更新一次参数效果上等价于 32。def train_one_epoch(model, dataloader, optimizer, scheduler, device, gradient_accumulation_steps1): model.train() total_loss 0 for step, batch in enumerate(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, labelslabels ) loss outputs.loss / gradient_accumulation_steps loss.backward() if (step 1) % gradient_accumulation_steps 0: torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step() optimizer.zero_grad() total_loss loss.item() * gradient_accumulation_steps return total_loss / len(dataloader)outputs.loss是模型内部计算好的交叉熵损失labels直接传给 model 就行不需要手动算 loss。clip_grad_norm_这句很重要BERT 这类大模型在训练后期很容易出现梯度爆炸不裁剪的话 loss 会突然变成 NaN然后整个训练就废了相当于没有后悔药。裁剪阈值写 1.0 是一个平衡选择太小会让收敛变慢太大等于没裁。训练过程中还需要固定随机种子否则每次跑出来的结果都不一样没法对比实验。utils.py里的 seed 函数一般会同时设置 Python、NumPy、PyTorch 和 CUDA 的随机种子这在复现实验时是刚需。4.4 评估与模型保存不看单点准确率训练完每个 epoch需要在验证集上做一次评估看的是整体的 precision、recall、F1而不只是准确率。尤其是文本分类场景类别不平衡是常态比如“售后”类样本占了 80%模型全都预测成“售后”准确率也有 80%但完全没用。from sklearn.metrics import classification_report, f1_score def evaluate(model, dataloader, device, id2label): model.eval() preds [] truths [] 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 batch_preds torch.argmax(logits, dim-1) preds.extend(batch_preds.cpu().tolist()) truths.extend(labels.cpu().tolist()) report classification_report(truths, preds, target_nameslist(id2label.values())) macro_f1 f1_score(truths, preds, averagemacro) return macro_f1, report评估里两个细节第一model.eval()和torch.no_grad()必须配对使用前者关掉 Dropout 和 BatchNorm 的训练行为后者关掉梯度计算两者都省显存缺一个都会出问题。第二averagemacro对每个类别的 F1 取算术平均能更公平地反映小类别的表现比weighted更严格。评估完把效果最好的模型保存下来model.save_pretrained(model_save/) tokenizer.save_pretrained(model_save/) import json with open(model_save/label2id.json, w, encodingutf-8) as f: json.dump(label2id, f, ensure_asciiFalse, indent2)这三个文件缺一不可加载模型时from_pretrained会读取整个目录label2id.json 则是预测时把数字索引翻译回中文类别名的关键。5. 避坑记录BERT 中文分类的几个经典翻车现场5.1 截断策略不当长文本分类全部偏向某一类现象训练 loss 正常下降验证集准确率也不算差但把真实场景的长文本丢进去预测结果几乎全成了同一类。原因默认 truncation 保留文本开头、丢掉结尾而很多关键结论恰恰出现在结尾。比如“整体不错但是售后…………”这种转折型文本模型只看到前半段正面内容自然分到“好评”。更麻烦的是训练数据大多短文本模型根本没学会处理截断后的长文本分布。解决把 tokenizer 的truncation_sideleft改成左截断让模型看到结尾更稳妥的做法是头部尾部各保留一段中间丢弃。我在代码里用的是右侧截断保留开头但实际业务文本如果确认是“结论在后”类型就要改成左截断这属于一行代码的改动tokenizer AutoTokenizer.from_pretrained(bert-base-chinese, truncation_sideleft)注意修改截断方向后训练和预测必须保持一致的截断方向否则模型看到的输入分布又对不上了。5.2 label2id 顺序随机导致预测结果张冠李戴现象训练了好几个 epoch验证集 F1 很高但拿新文本预测时模型输出的中文类别名称和文本内容明显对不上比如把“求开发票”分到了“投诉”。原因建立 label2id 映射时用了set(labels)去重集合是无序的每次运行生成的顺序都不一样。训练时映射关系写在 train.py 进程里没问题但保存模型时这份映射没有被固化预测时重新生成了一份顺序不同的映射表模型的数字索引和新映射表对不上。解决在data_processor.py里用有序方式生成映射表并保存到 json 文件labels sorted(set(all_labels), keylambda x: str(x)) label2id {label: idx for idx, label in enumerate(labels)} id2label {idx: label for label, idx in label2id.items()} with open(model_save/label2id.json, w, encodingutf-8) as f: json.dump(label2id, f, ensure_asciiFalse, indent2)核心原则label2id 的生成、保存、加载必须在三个环节里保持同一份文件不要依赖运行时内存里的映射关系。5.3 显存不够时乱调 batch size模型不收敛现象装下代码后一跑CUDA out of memory于是把 batch_size 从 16 调小到 4然后发现训练 loss 下降变得极慢模型效果一塌糊涂。原因batch size 变小后每个 step 的梯度噪声变大学习率不变的情况下模型参数更新方向非常不稳定。这其实是超参数联动的问题和模型本身没关系。解决batch size 减小的同时要么降低学习率比如从 2e-5 改成 1e-5要么用梯度累积补足有效 batch size。一个血泪经验是先确定单卡能承载的最大 batch size在这个基础上用梯度累积扩展最后再考虑调整学习率。5.4 一个通用的排查思路先过拟合一条数据现象训练集 loss 不降验证集指标也不行你分不清是代码 bug、数据问题还是模型架构问题。原因训练链路某个环节出了问题但报错信息不足以定位。解决每次启动正式训练前先只取 10 条训练数据把 batch size 也调成 10跑 5 个 epoch。如果 loss 能从 1.0 以上降到接近 0说明模型链路和反向传播没问题代码链路是通的问题出在数据规模或数据质量上。如果连 10 条数据都过拟合不了那就是代码有 bug没必要浪费一整个晚上的 GPU 时间在错误上。这个方法成了我每次跑新数据集的固定热身动作。6. 二次开发改造换数据、换标签体系、换任务类型的落地手法这套代码最值得称赞的地方是它的可替换性。换数据集只需要改 CSV 文件换分类类别数量只需要改num_labels自动读取映射表。如果你不止需要单标签分类想做多标签分类——比如一条用户反馈同时属于“产品质量”和“需要回访”——就把BertForSequenceClassification换成BertForMultiLabelClassification损失函数从 CrossEntropyLoss 换成 BCEWithLogitsLoss预测时从argmax改成阈值判断。outputs model(input_idsinput_ids, attention_maskattention_mask) logits outputs.logits probs torch.sigmoid(logits) # 多标签分类每个类别独立判断 preds (probs 0.5).int()如果文本特别长超过 512 个 token 上限常见做法是分段处理后做“投票聚合”每段独立预测概率然后按类别取平均。这种方案适合几千字的长文档分类场景。部署方面模型保存为model_save目录后可以原样用 PyTorch 的 TorchScript 序列化后走 C 推理也可以直接导成 ONNX 格式配合 ONNX Runtime 部署到 CPU 服务器上每秒处理几十条请求是足够的。最后分享一个已经固化成习惯的教训现在每接手一份新数据我都强制走一遍“10 条数据过拟合检查 - 全量训练 - 验证集分类报告逐类查看 - 抽 5 条真实预测结果人工复核”的流程缺一步就觉得心里不踏实。这份代码的价值就在于此——它把一个看似深不可测的 BERT 工程流程压缩成了几个可以逐行审查的文件希望你也能基于它快速调出一版属于自己的分类模型希望帮到你。本文还有配套的精品资源点击获取
锦
锦皓数字建站
深耕本土企业品牌数字化升级,专注原创端正雅致商务官网,从视觉设计到稳定运维全程保驾护航。