简介基于 pytorch-transformers 实现的 BERT 中文文本分类代码面向 NLP 入门与文本分类实践者也适合课程设计或论文复现。项目基于 THUCNews 20 万条新闻标题数据按 18 万、1 万、1 万划分为训练、验证与测试集覆盖财经、房产、股票、教育、科技、社会、时政、体育、游戏、娱乐 10 个类别代码包含文本预处理、模型加载、微调训练、评估与预测的完整流程。资源包共 29 个文件以 Python 脚本、txt 数据、json 与 xml 配置、预训练模型文件为主另有训练与评估结果图表、README 说明文档压缩包大小约 732.57MB。其中 Python 脚本负责数据预处理、模型训练与阶段评估txt 文件用于保存样本与输出结果模型权重便于直接加载使用。已有 3699 人学习该资源。借助内置的 BERT 模型权重与微调脚本使用者可直接复现实验也可替换为自己的中文数据集快速开展文本分类任务。 这两年做中文文本分类我最常被问到的一个问题是为什么从传统的词向量加分类器迁移到BERT这套方案早些年我用TF-IDF加LightGBM或者Word2Vec拼TextCNN说实话在小规模数据集上也够用。但一遇到长尾表达、一词多义、口语化文本传统方法就明显吃力。后来切到pytorch-transformers实现的BERT中文文本分类其实代码改造成本没有想象中那么高但对分类效果的提升是跨档次的。这篇文章就把这套完整代码拆开讲透包含数据预处理、模型搭建、训练、评估和预测以及我在实际项目中踩过的坑。这套代码适合谁有些读者刚刚入门还没搞清楚BERT的输入到底长什么样有些读者在用BERT做分类但训练过程总是过拟合或者loss不降还有一些是想把单条文本预测封装成接口的人。无论你属于哪一类下面这份代码骨架都能直接套用你只需要替换成自己的数据和标签就行。我个人更建议你一定要看懂每一步在干什么而不是复制粘贴跑完就算完事。1. 为什么选BERT做中文文本分类这套代码的设计目标先说结论BERT对中文文本分类最大的贡献是把语义理解这件事的门槛降了下来。传统做法要人工做大量特征工程字词权重、n-gram、句法特征都得考虑而且换个场景就要重新设计。BERT预训练模型在大规模中文语料上已经学出了很强的语言表征下游分类只需要在最后一层加一个分类器微调几个epoch就能达到一个很不错的baseline。这套代码的设计目标也很明确开箱即用且易于二次改造。整个流程是数据读取 - 文本编码 - DataLoader包装 - 模型加载 - 训练 - 验证 - 保存 - 预测。我给每个环节都保留了独立的函数或类这样你可以只替换数据读取部分或者只调整模型参数不至于牵一发动全身。另外一个容易被忽略的设计点是代码同时照顾了开发效率和部署需求。训练阶段用完整的训练循环保证可以观测每个batch的loss变化预测阶段则单独抽出一个函数输入一条文本直接返回标签和置信度。后续要接Flask、FastAPI还是单纯本地批量推理都不用再改核心逻辑。我用的依赖核心是pytorch-transformers这套接口。这里补充一句Hugging Face现在主推的库名是transformers老项目里也可能出现pytorch-transformers这个包名两者在BERT相关的常用接口上基本一致代码里按from transformers import xxx写就行第三方环境兼容性更好。我下面所有代码都以新版库为准。2. 环境搭建与依赖版本这块最容易出幺蛾子很多人跑BERT代码最大的障碍不在模型而在环境。pytorch-transformers相关的库迭代非常快不同版本之间API差异很大而且中文语料预训练模型的下载还需要网络条件稳定。我先给出一份我测试过能稳定运行的依赖组合。2.1 建议的依赖版本组合我用的是Python 3.8环境GPU环境下训练CPU环境也能跑但速度会慢很多。关键依赖如下torch 1.13.1 transformers 4.26.1 numpy pandas scikit-learn tqdmtransformers版本不建议追新有些新版API会调整例如BertForSequenceClassification的返回值结构就经历过变化。4.26.1这个版本对中文BERT兼容性很好网络上也容易找到对应文档。如果你是从老项目迁移看到pytorch_transformers这个导入名可以统一改成transformers。注意老版本有些模型名称的前缀可能带bert-下载的时候要写完整的模型id这个后面细说。2.2 安装命令与常见问题pip install torch1.13.1 transformers4.26.1 pandas scikit-learn tqdm如果安装慢可以用国内镜像把pip源指到清华或阿里云。但有一个坑我必须提醒不要在安装过程中随意升级torch的版本。有的依赖会强制要求更高版本的torch一升级torch.nn.DataParallel和一些算子行为可能变化之前调通的代码会无端报错。另外很多新手会遇到Loading bert-base-chinese时长时间卡住或直接报连接错误。这是因为Hugging Face模型托管在海外下载不稳定。我的做法是提前把模型文件用脚本下载到本地然后通过本地路径加载。下面这段代码可以先跑一次下载完成后会在本地生成缓存目录from transformers import BertTokenizer, BertForSequenceClassification model_name bert-base-chinese tokenizer BertTokenizer.from_pretrained(model_name) model BertForSequenceClassification.from_pretrained(model_name, num_labels2)如果你不想用Hugging Face官方源也可以去ModelScope等平台下载对应的二进制权重然后单独指定from_pretrained的路径。后面给的示例代码里我会把模型路径都抽成变量方便你替换。3. 中文文本预处理Tokenizer、Dataset与DataLoader的衔接BERT模型输入不像传统NLP直接把词语转换成id就行而是要经过tokenizer生成input_ids、attention_mask和token_type_ids三样东西。中文场景比英文稍微简单一点因为BERT的中文词典是基于字的不需要额外分词直接按字切分即可。但这里面的细节仍然不少。3.1 tokenizer到底做了什么以一句“这家酒店的早餐非常丰富”为例BertTokenizer会先做规范化再映射到词典中的id。同时会在句子开头加上[CLS]在句子末尾加上[SEP]。input_ids就是这些token对应的编号attention_mask标记哪些位置是真实的token哪些位置是padding补出来的token_type_ids在单句分类里全部为0如果是句子对任务才区分前后句。你需要设定一个max_length超过的截断不足的补[PAD]。比如定为64或128训练时所有样本统一长度才能组成一个batch。下面这段编码函数是整套代码的核心入口def encode_texts(texts, tokenizer, max_length128): encodings tokenizer( texts, truncationTrue, paddingTrue, max_lengthmax_length, return_tensorspt, ) return encodings注意truncationTrue和paddingTrue同时打开tokenizer会自动完成长度对齐。有一个容易踩的坑如果传入的是一个list则按batch处理如果传入的是单条字符串也要用[text]包一层否则返回的维度会不对后面拼DataLoader会直接报错。3.2 自定义Dataset类实际项目中数据通常存在CSV或Excel里包含文本列和标签列。下面这个Dataset封装可以直接使用from torch.utils.data import Dataset class TextClassificationDataset(Dataset): def __init__(self, texts, labels, tokenizer, max_length128): 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 int(self.labels[idx]) encoding self.tokenizer( text, truncationTrue, paddingmax_length, max_lengthself.max_length, return_tensorspt, ) return { input_ids: encoding[input_ids].squeeze(0), attention_mask: encoding[attention_mask].squeeze(0), labels: torch.tensor(label, dtypetorch.long), }注意这里padding方式用了max_length也就是说每一条样本都补齐到相同长度。这样取batch时就不需要额外做collate操作简单直接。缺点是会浪费一点显存但在128这个量级上完全可以接受。如果你的数据长短差异非常大可以考虑改成动态padding用collator按当前batch的最大长度补齐能省不少显存。3.3 标签与id的映射如果是二分类标签通常就是0和1如果是多分类建议做一个单独的映射文件或字典保存防止预测时标签错位。我习惯提前把标签统计出来并排序固定label_list sorted(set(train_df[label].tolist())) label2id {label: i for i, label in enumerate(label_list)} id2label {i: label for label, i in label2id.items()}这样做的好处是不管原始标签是中文还是数字代码里统一按整数id处理。训练结束后保存id2label预测阶段才能把模型输出的索引还原成可读标签。4. 模型核心搭建BERT分类模型并理解输出结构有了输入数据下一步就是构建模型。直接用transformers里的BertForSequenceClassification是最省事的它内部结构等于BERT编码器加一个分类头。但如果你想灵活调整分类头比如加Dropout层、做多任务学习那就要拆开写。4.1 直接用封装好的模型from transformers import BertForSequenceClassification model BertForSequenceClassification.from_pretrained( bert-base-chinese, num_labelslen(label2id), )这里num_labels会自动决定分类头的输出维度。模型前向传播时传入input_ids、attention_mask和labels返回一个SequenceClassifierOutput对象里面有loss和logits字段。训练阶段只需要把loss拿出来做反向传播推理阶段只需要用logits算预测值。4.2 手动拆分BERT编码器与分类头如果你的分类任务稍微复杂一点比如需要在BERT输出之上拼接其他特征或者做多标签分类那就得手动搭。下面这种写法把BERT编码器和分类头分开import torch.nn as nn from transformers import BertModel class BertClassifier(nn.Module): def __init__(self, pretrained_path, num_labels, dropout0.1): super().__init__() self.bert BertModel.from_pretrained(pretrained_path) self.dropout nn.Dropout(dropout) self.classifier nn.Linear(self.bert.config.hidden_size, num_labels) def forward(self, input_ids, attention_mask): outputs self.bert( input_idsinput_ids, attention_maskattention_mask, ) pooled outputs.pooler_output pooled self.dropout(pooled) logits self.classifier(pooled) return logits解释一下这里最关键的一个设计为什么拿pooler_output而不是last_hidden_state的均值池化pooler_output是BERT对[CLS]位置输出再接一层全连接和Tanh后的结果它是专门为分类任务设计的句子级表征。实际测试下来在大多数中文分类数据上它都优于简单均值池化。但也有例外如果你的输入是超长文本且截断了大量信息可能均值池化更稳定。我的建议是默认用pooler_output跑一版评估后可以再对比均值池化。4.3 类别不平衡问题与损失函数很多真实项目里类别是不均衡的比如投诉文本远少于正常文本。如果直接用自带的交叉熵损失少数类基本学不动。我的做法是传入class_weight给损失函数。封装模型时可以这样改class_weights torch.tensor([0.7, 1.3], devicedevice) loss_fct nn.CrossEntropyLoss(weightclass_weights)权重怎么算一个简单方案是用1 / 类别样本占比做归一化。但要注意权重不要设置得过于极端不然容易把模型带偏导致大多数样本都被预测成少数类。5. 训练策略与训练循环优化器、学习率调度和关键代码模型搭建完成进入训练阶段。BERT微调和其他深度学习模型训练有一些明显区别学习率通常很低训练轮数不需要太多还需要预热和衰减策略。这些细节直接影响收敛速度和最终效果。5.1 优化器与学习率BERT微调用的优化器是AdamW它是Adam的改进版主要修正了权重衰减的实现方式。学习率一般取2e-5到5e-5。我常用的是2e-5起步如果loss震荡明显降到1e-5。学习率调度上前10%的step做warmup让学习率从0缓慢升到设定值之后线性衰减到0。这个策略非常稳定基本不用过多调整。from transformers import AdamW, get_linear_schedule_with_warmup optimizer AdamW(model.parameters(), lr2e-5, weight_decay0.01) total_steps len(train_dataloader) * epochs warmup_steps int(total_steps * 0.1) scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_stepswarmup_steps, num_training_stepstotal_steps, )有读者会问不用官方分组的参数权重衰减行不行直接用AdamW对全部参数做weight_decay也行但更规范的做法是把LayerNorm和bias排除在权重衰减之外。transformers官方提供了一个get_parameter_names方法来自动识别这些参数不过在小规模任务上并不敏感。为了让代码简洁我这里没有做分组效果也没问题。5.2 训练循环代码下面是核心训练循环。考虑可读性我用最简单的写法每一步都打印batch lossdevice torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) epochs 4 for epoch in range(epochs): model.train() total_loss 0 for step, batch in enumerate(train_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 loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad() total_loss loss.item() if step % 50 0: print(fEpoch {epoch}, Step {step}, Loss {loss.item():.4f}) print(fEpoch {epoch}, Average Loss {total_loss / len(train_dataloader):.4f})这里有一个非常重要的细节梯度裁剪。clip_grad_norm_(model.parameters(), 1.0)这行看起来不起眼但BERT微调时偶尔会出现loss突然变成NaN的异常绝大多数情况都是梯度爆炸导致的。加上这行之后训练稳定性会明显提高。我一般把max norm设为1.0如果你发现梯度裁剪后收敛变慢可以试试1.5或2.0。5.3 显存不够怎么办梯度累积与混合精度如果你的GPU显存只有4G或6G可能会遇到CUDA out of memory。除了换更小的max_length最常用的方法是梯度累积。思路很简单不更新参数但反向传播累积梯度累计若干步后再更新一次。accumulation_steps 2 for step, batch in enumerate(train_dataloader): outputs model(...) loss outputs.loss / accumulation_steps loss.backward() if (step 1) % accumulation_steps 0: optimizer.step() scheduler.step() optimizer.zero_grad()注意我做了loss / accumulation_steps这一步很关键否则累积梯度的量级会变大等于隐式放大了学习率。如果你还嫌慢可以再用混合精度训练。在torch 1.13中可以直接用官方torch.cuda.ampGradScaler配合autocast代码加几行就能省差不多一半显存推理速度也有提升。但要注意混合精度在某些老显卡上支持不完整使用前确认算力在7.0以上。5.4 epoch数量与早停BERT微调不建议训练太多轮。中文分类任务我一般设3到5个epoch。在很多公开数据集上第2到第3个epoch效果就达到峰值了继续训练容易过拟合。更稳妥的办法是做早停每个epoch结束都在验证集上计算指标如果连续两个epoch指标没有提升就停止训练并回滚到最佳模型。6. 验证、评估指标与单条文本预测训练不能只看训练loss必须用验证集评估。验证集的构造也很重要尽量保持和真实应用场景一致。如果你拿到的数据有时间序列特征比如新闻按时间分布最好按时间切分而不是随机切分否则评估结果会虚高。6.1 验证循环from sklearn.metrics import accuracy_score, f1_score, classification_report def evaluate(model, dataloader, device): model.eval() all_preds [] all_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) logits model(input_idsinput_ids, attention_maskattention_mask) preds torch.argmax(logits, dim-1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.cpu().numpy()) acc accuracy_score(all_labels, all_preds) f1 f1_score(all_labels, all_preds, averageweighted) print(fAccuracy: {acc:.4f}, F1: {f1:.4f}) print(classification_report(all_labels, all_preds)) return acc, f1我很看重F1而不是只用准确率。特别是类别不均衡时准确率会被多数类拉高F1的weighted或macro版本更能反映少数类的表现。上面用的weighted会按样本量加权macro则直接取各类F1的平均。如果业务上更关注少数类的召回率你可以用recall_score单独观察。6.2 模型保存与加载训练结束模型和tokenizer都要保存。推荐方式如下save_path ./bert_model model.save_pretrained(save_path) tokenizer.save_pretrained(save_path)这里save_pretrained会同时保存权重文件和config.json。下次加载时model BertForSequenceClassification.from_pretrained(save_path) tokenizer BertTokenizer.from_pretrained(save_path)这样做的好处是加载时不需要再指定num_labelsconfig里已经记录了。不建议只保存state_dict因为还要手动维护config容易漏。如果你做了类别映射一定把id2label单独存一份json否则上线预测时只能看到数字不知道对应什么标签。6.3 单条文本预测函数实际部署的入口往往是单个字符串预测。下面这个函数可以直接复用def predict_single(text, model, tokenizer, id2label, max_length128, devicecpu): model.eval() encoding tokenizer( text, truncationTrue, paddingmax_length, max_lengthmax_length, return_tensorspt, ) input_ids encoding[input_ids].to(device) attention_mask encoding[attention_mask].to(device) with torch.no_grad(): logits model(input_idsinput_ids, attention_maskattention_mask) probs torch.softmax(logits, dim-1) pred_id torch.argmax(logits, dim-1).item() confidence probs.max().item() return id2label[pred_id], confidence返回的confidence是softmax后最大类别的概率这个值可以当作置信度使用。但我要提醒一点BERT在分布外样本上的置信度往往虚高实际部署时最好根据业务数据另设一个阈值比如小于0.8的样本打入人工审核队列。7. 复现这套代码时最容易踩的坑最后分享一些我从实际项目中总结的经验。这些细节不一定写在官方文档里但很影响最终效果和开发体验。7.1 中文文本里的特殊字符与空值中文数据里经常会有各种奇怪的字符比如全角空格、换行符、表情符号。BERT的tokenizer一般能处理但如果文本里混入了空值或非字符串内容比如数字、NaNtokenizer会直接报错或产生非法输入。我在预处理阶段一定会做两步一是把所有NaN填充成空字符串二是统一转成str类型。另外如果你的数据是爬下来的经常有大量HTML标签或URL建议在进模型前先清洗掉因为BERT的词典里没有URL这种词截断后可能占用大量token位置。7.2 max_length的选择不是越大越好很多初学者喜欢把max_length设成512认为覆盖全部信息更保险。但BERT的self-attention计算量是输入长度的平方量级设512显存占用和训练时间都会暴涨。实际经验是大部分中文短文本分类任务128已经够用超过200的样本占比往往很低。你可以先统计一下数据分布用95%分位长度作为max_length这样既不丢失太多信息又能节省计算资源。7.3 数据顺序与随机种子BERT训练对数据顺序不算特别敏感但如果不做shuffle连续batch的分布可能高度相似导致loss曲线波动很大。我在构建DataLoader时一定开启shuffleTrue。同时为了让实验可以复现我习惯在代码开头固定随机种子def set_seed(seed42): import random import numpy as np random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)这里多说一句torch.backends.cudnn.deterministic是否要开启开启后复现性更好但训练速度会下降一些。如果是做实验对比建议开启如果是跑生产模型速度优先可以不开。7.4 不同中文BERT模型的选择我默认用的是bert-base-chinese但针对不同业务场景可以替换成其他中文预训练模型比如hfl/chinese-roberta-wwm-ext、hfl/chinese-bert-wwm-ext。这类模型在全词掩码和更多中文语料上做了优化分类效果通常比原始BERT好1到2个点。换模型时只需要修改model_name变量代码其他部分完全不用动。这也是Hugging Face这套接口设计得好的地方。唯一要注意的是下载模型前确认磁盘空间足够这类模型文件一般在400MB左右。7.5 验证集和测试集的区分如果你的目标是发论文或者做严谨的对比一定要把数据集切成训练集、验证集、测试集三部分。验证集用来做模型选择和调参测试集只能在最终评估时用一次。很多人在实际项目中图省事只用一份数据又训练又验证最后模型看起来指标很高上线遇到新数据就现原形。我的习惯是先按标签做分层抽样切分比如训练集80%、验证集10%、测试集10%并固定随机种子。这样每次实验的数据划分一致调参前后对比才有意义。如果你只是想快速跑通一个闭环那训练集加验证集两份就够。上面给的代码已经覆盖了完整流程你只需要读入数据调用TextClassificationDataset和evaluate函数即可。后面的优化方向可以是做交叉验证、加入类别权重、换更好的预训练模型、甚至用对抗训练增强泛化能力。但先把基础链路跑通再谈优化这是我最想强调的一点。这套代码我自己在实际项目中反复用过很多次每次换业务场景只需要修改数据读取和标签映射两部分。你如果现在就有一份中文分类数据在手完全可以打开编辑器按上面步骤一步步写完跑通一个epoch看看loss趋势。刚开始精度不高很正常BERT微调真正稳住效果一般都在第二个epoch之后不要第一轮看到loss不降就急着改代码。最后再提醒一句保存好你的id2label否则训练完成后连预测结果都读不懂。本文还有配套的精品资源点击获取