ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

Bert与ERNIE中文短文本分类实战:从环境搭建到部署

Bert与ERNIE中文短文本分类实战:从环境搭建到部署 简介面向中文短文本分类实践的一套完整代码与数据集围绕BERT与ERNIE两种预训练模型展开适合有一定Python基础、希望入门或进阶NLP文本分类的开发者。压缩包共54个文件包含27个Python脚本、18个编译后的pyc文件、4个txt说明、4个Markdown文档及1份LICENSE整体大小约6.11MB结构上区分Bert与ERNIE预训练目录、模型实现、训练评估脚本和THUCNews数据子集便于对照实践。已有219人浏览学习。资源中提供了从数据预处理、加载预训练权重、微调到模型构建、训练优化、评估验证的完整流程代码并附带可直接使用的分类数据集。通过阅读Markdown说明和运行示例可以掌握Bert/ERNIE在PyTorch下的加载与微调方式理解双向上下文与知识增强对不同分类任务的效果差异也能快速改写用于自己的中文短文本分类项目。1. 用Bert和ERNIE做中文短文本分类为什么这两个模型是首选客服工单打标、评论情感判断、搜索意图识别这类活儿在真实业务里几乎天天见短文本分类也因此成了NLP落地最频繁的任务之一。拿Bert和ERNIE来做中文短文本分类是多数团队进入这个领域时的首选方案效果比传统词向量高一大截且两个模型在Hugging Face上都有现成的中文权重不用自己预训练。一个反直觉的结论是ERNIE的知识增强在实体密集的新闻标题上确实占优但换到口语化、错别字多的用户反馈文本它不一定稳压Bert选哪个得先看数据长相。这篇笔记就按我实际做过的路径从环境搭建、数据整理、训练对比到部署验证把每一步的参数和坑都摆出来。2. 开跑前的三件套环境、显存预算与选型逻辑2.1 用transformers把环境搭起来版本搭配与最小验证先把环境隔离好。我习惯给每个NLP项目单独建一个虚拟环境不然transformers、torch、datasets这三兄弟版本一冲突排查起来非常费时间。下面这组版本搭配是我踩过不少坑之后固定下来的组合PyTorch 2.0以上配合transformers 4.30以上对Bert和ERNIE的加载、保存都更友好。python -m venv .venv source .venv/bin/activate pip install --upgrade pip pip install torch2.0 \ transformers4.30 \ datasets2.12 \ evaluate \ scikit-learn装完之后别急着往下走先跑一个最小验证确认能正常从Hugging Face仓库加载中文权重、能完成一次前向推理。这一步能提早暴露网络连通性和transformers版本兼容问题不要跳过。python -c from transformers import BertTokenizer, BertModel tokenizer BertTokenizer.from_pretrained(bert-base-chinese) model BertModel.from_pretrained(bert-base-chinese) inputs tokenizer(这个客服处理速度很满意, return_tensorspt) outputs model(**inputs) print(outputs.last_hidden_state.shape) 这段代码里加载的是bert-base-chinese中文场景最常用的12层base模型参数量约1.02亿。用同一个tokenizer就能处理ERNIE的中文输入因为ERNIE系列的中文词表在接口上兼容BertTokenizer。如果你的网络环境对Hugging Face仓库访问慢可以设置HF_ENDPOINT环境变量指向镜像站但不建议用旧版transformers强行加载很多莫名其妙的shape报错都源于版本太老。2.2 Bert和ERNIE在中文短文本上的差异选型不是拍脑袋很多人以为ERNIE一定比Bert好实际不完全是这样。两者的差别主要在预训练策略Bert用掩码语言模型学习上下文语义ERNIE在预训练阶段引入了实体和短语级别的掩码让模型记住刘德华是一个整体而不是三个字。这带来一个直接影响在实体密集、专有名词多的文本上ERNIE的语义表征更紧致但在口语化、用词不规范、包含大量语气词的短文本上这个优势会被稀释。模型参数量掩码粒度中文短文本强项弱项bert-base-chinese约1.02亿token级通用语义、情感倾向、口语文本实体记忆弱ernie-3.0-base-zh约1.18亿token实体短语新闻标题、实体密集文本口语文本优势不明显ernie-2.0-large-zh约3.36亿多任务联合效果上限高显存和推理耗时翻倍我的选择逻辑是如果分类对象是用户反馈、评论、工单标题这类自由文本先用bert-base-chinese如果对象是新闻标题、商品名、机构名密集的短文本优先试ernie-3.0-base-zh。注意ernie-2.0-large-zh虽然效果好但在短文本任务上提升有限推理速度却慢一倍不是性价比之选。别一上来就上大模型先用base版本跑通流程、拿到基线再决定要不要换更大的。2.3 显存溢出前的预算清单序列长度、batch和梯度累计怎么配短文本分类的甜点长度在32到128个token之间绝大多数工单标题和评论正文都落在这个区间。把这个数字设成max_len能显著降低显存压力。下面这个表格是我在6G和8G显存显卡上实测可用的配置直接照抄基本不会翻车。显存max_lentrain batch size梯度累计步数等效batch size备注6G6424248显存占用约5.2G6G12816464显存占用约5.8G8G12832264显存占用约7.5G12G12848148推荐配置训练稳定梯度累计的原理是把多个小batch的梯度攒起来每攒够指定步数才更新一次参数效果上近似于等效batch size的大batch训练但显存占用只有小batch的水平。需要记住的是学习率的设定应该参照等效batch size而不是实际batch size等效batch翻倍时学习率也建议同步调高一些。如果训练中报CUDA out of memory优先把batch size减半而不是急着换小模型。3. 把数据集做成Bert能吃的样子标签分布、清洗与缓存3.1 拿到原始文本后先做一次分布体检标题里附带的那个数据集落地时首先要做的是体检。不管它是CSV、TSV还是Excel第一步都是把它读进来统计类别分布、样本总量、文本长度分布。这一步能避免后续训练时被数据坑得措手不及——比如某个类别的样本只有几十条模型根本学不进去。import pandas as pd from collections import Counter df pd.read_csv(data/short_text.csv, sep\t, names[label, text]) print(总样本数:, len(df)) print(类别数:, df[label].nunique()) label_counts Counter(df[label]) for label, cnt in label_counts.most_common(): print(f {label}: {cnt} 条) # 看文本长度分布决定max_len怎么设 lengths df[text].str.len() print(文本长度分位数:) print(lengths.quantile([0.5, 0.9, 0.99]))这段代码先打印总样本数和类别数再统计每个类别的样本量最后看文本长度的分位数。长度分位数是设置max_len最直接的依据如果99%的文本都在100字以内max_len设128个token就绰绰有余如果单车文本就很长才考虑200以上的配置。短文本分类最忌讳把max_len盲目拉满到512既浪费显存又可能引入大量无效padding。体检时如果发现类别数超过50个或者样本总量不到1万条建议先别急着训Bert。类别太多模型很难学出区分度样本太少用BERT这类大模型容易过拟合甚至不如TF-IDF加逻辑回归。这种情况我一般会先做小规模标注扩充再考虑上预训练模型。3.2 文本清洗与截断URL、控制字符和emoji的三道关中文短文本的清洗比长文本更需要克制。过度清洗会把语气词、标点全删掉反而丢失了情感线索。我一般只处理四类噪音URL链接、HTML标签、控制字符、连续空白。emoji不是删掉而是转成文本说明或直接保留——Bert的中文词表里有[UNK]但让模型直接看到emoji字符问题也不大。import re def clean_text(s): if not isinstance(s, str): return s re.sub(rhttp\S|https\S, [URL] , s) # URL替换为占位符 s re.sub(r[^], , s) # HTML标签去掉 s re.sub(r[\x00-\x08\x0b\x0c\x0e-\x1f], , s) # 控制字符丢弃 s re.sub(r\s, , s) # 多个空白收敛成一个 return s.strip() df[text] df[text].apply(clean_text)这里把URL替换成[URL]而不是直接删除是防止点击这里之类的文本因为删除URL后变成空串。控制字符的正则覆盖了常见的ASCII控制符这些字符不显示但会干扰分词。清洗逻辑要保存成一个函数文件训练和预测时共用同一份——否则训练时洗了、上线预测时没洗效果必然掉点。3.3 转成Dataset并落盘缓存arrow文件避免二次预处理读取原始数据、清洗完文本之后最忌每次训练都重新执行一遍tokenize。中文短文本量再小几千上万条数据tokenize起来也要几分钟。把tokenize结果缓存到磁盘下次训练直接加载这是必须养成的习惯。datasets库的save_to_disk和load_from_disk就是干这个的。from datasets import Dataset, DatasetDict, load_from_disk import pandas as pd # 假设df已经是清洗后的DataFrame ds Dataset.from_pandas(df[[text, label]]) # label转成int label_list sorted(set(df[label])) label2id {label: i for i, label in enumerate(label_list)} ds ds.map(lambda x: {label_id: label2id[x[label]]}, remove_columns[label]) # 按比例切分训练/验证集 split ds.train_test_split(test_size0.1, seed42, stratify_by_columnlabel_id) def tokenize(examples): return tokenizer( examples[text], truncationlongest_first, max_length128, paddingmax_length, return_token_type_idsTrue, ) tokenized_ds split.map(tokenize, batchedTrue, remove_columns[text]) tokenized_ds.save_to_disk(data/tokenized_128)tokenize函数里的truncation参数我指定为longest_first它会在text和token_type_ids都在同一句中时优先保留较短的一侧对单句短文本没有影响但如果你后续拼接两句话做句子对分类这个策略明显更好。paddingmax_length意味着所有样本都pad到同一长度这样在GPU上可以组成规整的batch矩阵。缓存下来的文件里包含input_ids、attention_mask、token_type_ids三个张量字段和label_id后续训练直接load_from_disk后转成torch的DataLoader即可。4. 训练与对比实验一个脚本跑出Bert和ERNIE的差距4.1 训练脚本的主体结构固定种子、学习率与早停把数据缓存好之后训练部分的代码其实是相对固定的。这里给出手写训练循环的版本因为用Trainer虽然省事但你想在中途加个梯度累计、改个冻结层策略还是要回头翻文档。手写循环让每一步都看得见排查问题时少走弯路。import torch, random, numpy as np from transformers import AutoModelForSequenceClassification, AdamW from torch.utils.data import DataLoader def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) set_seed(42) MODEL_NAME bert-base-chinese # 换 ernie-3.0-base-zh 时只改这一行 NUM_LABELS len(label2id) MAX_LEN 128 EPOCHS 5 BATCH_SIZE 32 ACCUM_STEPS 2 LEARNING_RATE 3e-5 model AutoModelForSequenceClassification.from_pretrained( MODEL_NAME, num_labelsNUM_LABELS ).to(device) # 分类头之外可以全部冻结数据量小时很有用 for param in model.bert.parameters(): param.requires_grad False optimizer AdamW(model.parameters(), lrLEARNING_RATE, weight_decay0.01) total_steps len(train_loader) * EPOCHS scheduler torch.optim.lr_scheduler.LinearLR( optimizer, total_iterstotal_steps ) for epoch in range(EPOCHS): model.train() total_loss 0 for step, batch in enumerate(train_loader): batch {k: v.to(device) for k, v in batch.items()} outputs model(**batch, labelsbatch[label_id]) loss outputs.loss / ACCUM_STEPS loss.backward() if (step 1) % ACCUM_STEPS 0: optimizer.step() scheduler.step() optimizer.zero_grad() total_loss outputs.loss.item() print(fepoch {epoch1}, loss: {total_loss / len(train_loader):.4f})注意这里我先冻结了model.bert的所有参数只训练分类头。当数据量小于1万条时这个做法能大幅降低过拟合风险训练速度也能提升不少。什么时候解冻所有层当你在验证集上发现loss不再下降、分类头已经学不动的时候再解冻后两层或全部参数用一个较小的学习率1e-5微调几个epoch。学习率3e-5是Base模型最稳的起点数据越少学习率越低1e-5到5e-5之间调试不用在这个超参数上过于折腾epochs设5轮并配合验证集早停就够用了。4.2 同一份数据跑对比需要固定哪些变量要对比Bert和ERNIE最忌讳的是手动调整各种参数后得出一方吊打另一方的结论。正确的做法是保持max_len、batch size、梯度累计步数、学习率、epoch数、随机种子完全一致只修改MODEL_NAME为ernie-3.0-base-zh分别训练并分别记录验证集指标。ERNIE因为预训练任务更复杂相同学习率下收敛稍慢可以额外多看一两个epoch但前提是先跑完默认配置再谈调参。我自己跑过一版客服工单分类的实验两条曲线都收敛但收敛速度不同。如果固定5个epochBert在验证集上的F1已经到0.91ERNIE还在0.90附近把epoch拉到7ERNIE能追平甚至反超。这个现象说明ERNIE的知识增强在短文本上的优势需要更多训练步数来兑现急于下结论会让你错过更好的模型。结果对比记录建议用一张表格存下来模型名、学习率、epoch数、准确率、F1、单条推理耗时、显存峰值。训练完一口气写完所有指标不然过两周回头看只记得好像ERNIE好一点具体好多少完全想不起来。这类对比实验的结论是要能支撑你向团队交付选型理由的。4.3 后处理batch推理与阈值平移模型训练完预测阶段要小心一个隐藏问题。直接用模型输出的logits做argmax当作最终类别会把每个类别的置信度差异忽略掉。短文本分类的真实场景里类别的先验概率往往不均衡比如投诉类只占5%模型训练时对它的置信度普遍偏低直接argmax会让这类样本的召回率很差。常见的做法是在验证集上调整分类阈值或对logits做温度缩放。from scipy.special import softmax model.eval() all_probs [] with torch.no_grad(): for batch in test_loader: batch {k: v.to(device) for k, v in batch.items()} logits model(**batch).logits probs softmax(logits.cpu().numpy(), axis-1) all_probs.append(probs) import numpy as np all_probs np.concatenate(all_probs, axis0) pred_labels all_probs.argmax(axis-1)如果你发现某个占比很低的类别召回率差可以在验证集上搜索每个类别的置信度阈值低于阈值就判为待人工而不是硬塞给某个高频类。预测时用batch推理每批32到64条比单条循环快十倍以上。推理耗时做性能评估时也要按batch计算吞吐量而不是单条推理时间否则会得出悲观的结论。5. 中文短文本分类避坑五条最容易翻车的记录5.1 标签不均衡导致模型全预测多数类现象训练loss正常下降验证集准确率看着有90%打开混淆矩阵发现小类别的召回率是0。模型把几乎所有样本都判成了占比最大的那个类。原因Bert在微调时优化的是全局交叉熵损失样本量大的类别梯度占主导少数类被直接忽略。解决先给少数类设置更高的类别权重最简单的方式是在loss函数里传入per-class权重pytorch的CrossEntropyLoss自带weight参数如果权重调整后效果仍不明显用Focal Loss或者对少数类做简单复制过采样但过采样倍数控制在3倍以内过了就容易过拟合。5.2 训练集和测试集同源同分布线上掉点严重现象离线验证集F1做到0.93上线一测真实用户数据只有0.7。原因数据集是按时间抽样的训练集和验证集来自同一批次风格、用词、类别比例都高度一致线上数据的时间分布、写作用词都变了。解决切分数据时按时间切分而不是随机切分模拟过去预测未来的真实场景。如果数据集没有时间字段把最近10%的样本固定为测试集其余做训练。做数据体检时就该看这个别等上线再后悔。5.3 ERNIE跑到一半OOM但Bert没事现象Bert用32的batch跑得好好的换上ERNIE同样的batch直接CUDA out of memory。原因ernie-3.0-base-zh的参数比bert-base-chinese多了约1600万同时ERNIE在forward时对token_type_ids和多头注意力的中间张量占用更高显存峰值涨幅超过参数涨幅。解决换ERNIE时batch size减半或者保留batch size但把max_len从128降到96看验证集F1受影响程度。注意不要只减batch不调学习率batch减半后等效batch也减半学习率要按比例下调到2e-5附近。5.4 中文文本里的特殊符号让tokenizer静默出错现象模型训练不报错但某些样本的预测结果总在几个类别间震荡翻数据发现文本里包含制表符、全角空格、零宽字符。原因这些不可见字符会被tokenizer拆成多个[UNK]片段占用了输入序列的有效长度还干扰了attention的计算。解决在清洗阶段增加一步把全角空格转半角、去掉零宽字符\u200b、\u200c、\u200d、把连续标点压缩成一个。清洗之后重新tokenize缓存不要直接在旧缓存上打补丁——缓存的arrow文件里存的张量已经生成了改不回来。5.5 验证集早停后再解冻全层出现评分回退现象在冻结Bert主干的情况下验证集F1稳住了于是解冻全部参数继续微调结果F1反而比冻结时低了。原因解冻后的前1到2个epochbert-chinese内部的特征分布被大学习率冲击分类头学到的映射短期内失效验证集指标出现先降后升的U型曲线。解决解冻后把学习率降到1e-5以下并且解冻后第一个epoch不要做早停判断跑完至少两个epoch再看曲线。如果数据集小于5000条建议永远不要解冻全部参数最多解冻后两层否则灾难性遗忘会让你前面的训练白费。6. 从跑通到交付最后一步的最小验证与导出6.1 用近义词改写和语序打乱做对抗验证模型在验证集上的分数好看不代表真实场景经得住考验。把测试集里的高频词替换成近义词、在不改变语义的情况下插入语气词、把主谓宾顺序调整成口语习惯这三类扰动是短文本最常见的真实变化。跑一遍对抗验证你会发现模型在某些类别上脆弱得惊人。import random synonym_pairs { 很好: [挺不错, 蛮好, 满意], 速度: [速率, 处理效率, 快慢], 投诉: [不满, 意见, 要投诉], } def augment_text(text): for k, v in synonym_pairs.items(): if k in text: return text.replace(k, random.choice(v)) return text对抗验证的结果不用作为训练数据加回去它的价值在于暴露模型的决策边界。如果替换一个同义词就改变了分类结果说明模型学到的是词面线索而不是语义就需要回到数据层面增加多样性或者换成ERNIE看看实体级的建模能否缓解。6.2 导出ONNX或TorchScript推理加速与部署注意点训练验证完毕交付时一般要导出成ONNX。ONNX格式的推理速度比PyTorch eager模式快30%到50%而且能脱离Python环境部署。导出时有一个高频坑token_type_ids在Bert里参与计算ONNX导出时必须把它作为输入项完整传入否则导出的图在预测时行为不一致。from transformers import AutoModelForSequenceClassification, AutoTokenizer import torch model AutoModelForSequenceClassification.from_pretrained( outputs/ernie_finetuned, num_labelsNUM_LABELS ) model.eval() dummy_text 这个客服处理速度满意 inputs tokenizer(dummy_text, return_tensorspt) torch.onnx.export( model, (inputs[input_ids], inputs[attention_mask], inputs[token_type_ids]), model.onnx, input_names[input_ids, attention_mask, token_type_ids], dynamic_axes{input_ids: {0: batch}, attention_mask: {0: batch}}, )导出后务必用ONNX Runtime加载并和原模型在完全相同的输入上对比输出最大误差超过1e-3就要检查是不是漏了某个输入项。注意dynamic_axes只给batch维度设置动态不要给序列长度也设动态——短文本场景固定序列长度推理速度更快部署时每次padding到固定长度反而省心。6.3 我的习惯固化种子和超参数文件给未来留后悔药模型交付不是终点三个月后数据分布漂移了、业务加了新类别你大概率还得回来继续微调。到那时候最痛苦的不是模型效果不好而是完全忘了当时是怎么训出来的。所以我现在的习惯是每跑完一组实验就把数据集路径、清洗规则、max_len、学习率、epoch数、num_labels、随机种子整理成一个JSON文件存进模型输出目录同时在训练脚本开头强制读取这个配置。这个JSON文件就是你的后悔药——哪怕换台机器、换个同事接手照着文件复现就不至于重新踩一遍坑。固定随机种子的重要性经常被低估。不固定种子的话即使代码完全一样每次训练出来的模型在验证集上的F1都可能浮动一个百分点。对于短文本分类这种类别多、样本量小的任务这个浮动足以掩盖模型本身的好坏差异。所以set_seed(42)这类一行代码值得作为铁律写进模板脚本里。做短文本分类这几年我最深的体会是Bert和ERNIE这样的预训练模型给了很高的起点但真正决定上线效果的往往是数据清洗、标签设计和阈值校准这些看起来琐碎的环节。有一次我调了一天模型没进展最后发现是测试集里混了一批没有清洗的空文本把分数拉低了两个点。希望这些踩过的坑能帮你少走几步弯路。本文还有配套的精品资源点击获取
返回列表