FEATURED · 精选文章

PyTorch从零实现BERT文本分类:完整指南与踩坑记录

发布时间 / 2026/9/8 8:06:03
来源 / 创域科博编辑部
栏目 / 资讯中心
PyTorch从零实现BERT文本分类:完整指南与踩坑记录 简介面向 PyTorch 与自然语言处理开发者这套代码资源完整实现了基于谷歌 BERT 模型的自然语言处理流程侧重从理论到工程的过渡环节可帮助读者快速理解双向 Transformer 结构、预训练与微调思想并顺利跑通文本编码、模型加载、下游任务训练和预测推理。项目围绕模型定义、训练器、数据集和主程序组织共 29 个文件其中 25 个 Python 脚本是核心另有 HTML 说明文件、TXT 文本说明、Git 忽略与属性配置等压缩包仅 18KB轻量便于下载、阅读和二次修改。目前已有 1544 人学习浏览受到 NLP 初学者和算法工程师的关注。通过研读这份代码可以掌握 BertTokenizer 的序列预处理、BertModel 的调用方式、注意力掩码的使用以及基于交叉熵损失和 AdamW 优化器的微调训练框架同时还能参考目录内脚本分工学习如何搭建一个简洁可扩展的 PyTorch NLP 项目为后续解决问答、情感分析、序列标注等实际任务提供基础。 我去年接了一个文本分类的项目需求是把一批客服工单自动打上标签。调研了一圈最终选了BERT作为核心模型。网上关于BERT的教程多如牛毛但真正能照着从头到尾把代码写出来、把模型训起来、还能处理各种工程细节的资料少之又少。很多文章要么直接调HuggingFace一行代码加载模型要么就是贴一段残缺不全的代码让你自己去猜。这篇博文我想分享的是基于PyTorch从数据处理、模型搭建到训练推理完整走一遍BERT的实现思路和踩坑记录不依赖HuggingFace封装的情况下怎么把BERT用在自己的项目里。1. 环境准备与依赖版本策略这一关其实就能劝退一半人先说环境。BERT这类Transformer模型对PyTorch版本、CUDA版本、Python版本极其敏感很多项目跑不起来根本不是代码问题而是版本不匹配。我建议在动手写代码之前先把环境彻底固定下来。1.1 PyTorch安装与CUDA匹配我的环境是Python 3.10.11 PyTorch 2.8.0 CUDA 12.1的组合包实际测试下来非常稳定。深度学习框架不要装最新版本也不要装太老的版本选一个经过了大量项目验证的稳定组合是最省心的。PyTorch 2.x系列在性能上有明显提升尤其是torch.compile功能但这个功能在BERT这种中小规模模型上收益有限反而可能带来编译时间开销所以建议先关掉跑通基础流程再加。安装PyTorch的时候很多人会栽在CUDA的坑里。这个坑有几个典型表现_torch.cuda.is_available()_返回False、运行时报错找不到CUDA驱动、或者CUDA和cuDNN版本不配套。如果是在Linux服务器上装我建议用Miniconda管理环境这里分享我常用的创建命令conda create -n bert_env python3.10 conda activate bert_env pip install torch2.8.0 torchvision0.19.0 --index-url https://download.pytorch.org/whl/cu121安装完以后用下面这段代码验证环境是否可用import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU mode)如果输出里_cuda.is_available()_是False优先检查显卡驱动而不是重装PyTorch。大版本匹配原则是这样的驱动支持的最高CUDA版本必须高于PyTorch所需的CUDA版本。比如PyTorch 2.8.0需要CUDA 12.1那驱动版本至少要支持CUDA 12.1及以上。Windows系统还要注意显卡是否支持CUDA——有些轻薄本的MX系列显卡算力很低跑BERT这种模型哪怕能装上CUDA训练速度也慢得让人怀疑人生。1.2 其它依赖库的选择BERT实现还需要Numpy、Pandas用于数据处理Matplotlib用于画loss曲线Tqdm显示训练进度条。我不建议在项目初期就引入过多的数据处理库很多操作Numpy原生就支持。等到数据量真的到了百万级再考虑引入更重的工具也不迟。另外_transformers_这个库我建议还是装上虽然后面我们会自己实现模型但分词器Tokenizer用的是HuggingFace的BertTokenizer这一点后面细说。注意即使最终目标是自己实现BERT模型也强烈建议安装transformers库。原因在于预训练模型的权重文件、词表文件都需要从HuggingFace下载而且你要用它的分词器来对齐输入格式。完全从零训练一个BERT模型的成本极高需要大量的数据和算力绝大多数场景下的合理做法是加载Google发布的预训练权重再做微调。2. BERT核心机制拆解写代码前必须想清楚的三件事写BERT代码之前有几个概念必须理解透彻否则代码写出来也是稀里糊涂的。2.1 BERT到底在做什么BERT的全称是Bidirectional Encoder Representations from Transformers核心思想是用Transformer的Encoder部分对文本进行双向编码。所谓的双向指的是模型在理解一个词的时候同时结合它左边和右边的上下文信息。这和传统的语言模型只从左往右读有本质区别。BERT在预训练阶段用了两个任务一个是掩码语言模型Masked Language ModelMLM随机遮盖输入中15%的词让模型根据上下文预测被遮住的词另一个是下一句预测Next Sentence PredictionNSP判断两个句子是否是连续的上下文。这两个任务决定了BERT输入格式和输出头的设计也直接关系到训练代码怎么写。2.2 输入到底长什么样这是新手最容易绕晕的地方。一句话要进入BERT需要被转换成三个向量序列输入类型作用示例input_ids每个token在词表中的索引[CLS], 我, 爱, 自然, 语言, [SEP] → [101, 2769, 4263, 6335, 6959, 102]token_type_ids区分句子A和句子Bsegment embedding第一个句子全部为0第二个句子全部为1attention_mask标记哪些是真正的文本哪些是padding真实token为1padding为0这三个向量是BERT所有代码的输入基础_input_ids_决定模型读什么_token_type_ids_告诉模型句子边界在哪_attention_mask_告诉模型哪些位置是应该被忽略的填充内容。2.3 为什么需要[CLS]和[SEP]标记[CLS]标记放在每个输入序列的开头它的最终隐藏状态被当作整个序列的语义表示用于分类任务。可以把[CLS]理解为一个信息汇聚点经过多层Transformer编码后这个位置的向量包含了整个句子的综合信息。在预训练阶段[SEP]标记用来分隔两个句子同时让模型知道句子之间的边界。提示如果你只用一个句子做分类比如情感分析输入就是[CLS] 句子token [SEP]。如果是两个句子做匹配比如问答、文本蕴含输入就是[CLS] 句A [SEP] 句B [SEP]。3. 数据预处理把文本变成BERT认识的数字数据预处理是整个流程中最枯燥、但也是最容易出错的部分。很多模型的性能上不去问题并不出在模型结构上而是预处理阶段就对不齐tokenizer的词表和模型预训练时的词表不一致、padding策略不对、attention_mask忘记加诸如此类。3.1 分词与词表对齐BERT使用WordPiece分词算法它会把unhappiness拆成[un, happiness]或者[un, happiness]具体拆分方式取决于词表。所以不能用自己训练的jieba之类的分词器去分必须使用和预训练权重匹配的BertTokenizer来加载词表并分词from transformers import BertTokenizer tokenizer BertTokenizer.from_pretrained(bert-base-chinese) text 我爱自然语言处理 result tokenizer(text, max_length128, paddingmax_length, truncationTrue, return_tensorspt) print(result[input_ids]) print(result[attention_mask]) print(result[token_type_ids])这里的一行代码背后做了很多事_paddingmax_length_会把所有文本都补齐到128的长度_truncationTrue_会把超过128的文本截断_return_tensorspt_表示返回PyTorch的张量格式。3.2 手写一个完整的Dataset类实际工程项目里数据往往存在CSV或者Excel文件里需要自己写Dataset类来管理数据和标签的加载。以情感二分类为例标签0表示负面1表示正面import torch from torch.utils.data import Dataset import pandas as pd class TextClassificationDataset(Dataset): def __init__(self, file_path, tokenizer, max_length128): self.data pd.read_csv(file_path) self.tokenizer tokenizer self.max_length max_length def __len__(self): return len(self.data) def __getitem__(self, idx): text str(self.data.iloc[idx][text]) label int(self.data.iloc[idx][label]) encoding self.tokenizer( text, max_lengthself.max_length, paddingmax_length, truncationTrue, return_tensorspt ) return { input_ids: encoding[input_ids].squeeze(0), attention_mask: encoding[attention_mask].squeeze(0), token_type_ids: encoding[token_type_ids].squeeze(0), label: torch.tensor(label, dtypetorch.long) }这里有几个细节值得注意。_squeeze(0)_是因为tokenizer默认会返回一个batch维度形状是[1, seq_len]取出来的时候需要把batch维度去掉变成[seq_len]的形式这样DataLoader在组合batch的时候才能正确堆叠成[batch_size, seq_len]。另一个细节是__getitem__里不要做任何数据增强或复杂的重处理保持轻量是Dataset设计的原则重活应该提前处理好或者放到collate_fn里。3.3 手写HuggingFace数据预处理上面直接用了HuggingFace的tokenizer虽然底层是封装好的但理解它的实现原理还是有必要的。_wordpiece分词的核心流程如下先做基本清理小写化、去掉特殊字符、按空格切分中文则按字切分对每个词判断是否能直接匹配词表如果不能尝试从最长的子串开始匹配##开头的token表示这个是某个词中间的部分如果连一个字的子串都匹配不上标记为[UNK]未知词理解了这个过程就会发现为什么中英文的处理方式不同英文有明确的空格分隔而中文每个字都需要被切出来。这也是为什么中文BERT要有专门的Chinese词表。3.4 数据的划分与保存预处理完之后的数据划分也很关键。我习惯用_sklearn_的_train_test_split_来划分训练集和验证集保证标签的类别分布与原始数据一致from sklearn.model_selection import train_test_split train_texts, val_texts, train_labels, val_labels train_test_split( texts, labels, test_size0.1, random_state42, stratifylabels )_stratifylabels_这个参数在类别不平衡的情况下尤其重要。如果不加可能验证集里没有某个类别的样本导致验证指标失真。4. 模型主体搭建从Embedding到Transformer Encoder现在终于到了代码实现的核心部分。我会从零实现一个基础的BERT模型覆盖Embedding层、多头自注意力层、前馈网络层以及最终分类头。这样做的价值在于你会真正理解BERT的每个模块做了什么而不仅仅是调用别人封装好的接口。4.1 Embedding层实现BERT的输入向量由三部分相加得到token embedding segment embedding position embedding。代码实现如下import torch import torch.nn as nn class BERTEmbedding(nn.Module): def __init__(self, vocab_size, hidden_size, max_length, dropout0.1): super().__init__() self.token_embedding nn.Embedding(vocab_size, hidden_size) self.segment_embedding nn.Embedding(2, hidden_size) self.position_embedding nn.Embedding(max_length, hidden_size) self.layer_norm nn.LayerNorm(hidden_size) self.dropout nn.Dropout(dropout) def forward(self, input_ids, token_type_ids): seq_length input_ids.shape[1] position_ids torch.arange(seq_length, deviceinput_ids.device).unsqueeze(0).expand_as(input_ids) token_emb self.token_embedding(input_ids) segment_emb self.segment_embedding(token_type_ids) position_emb self.position_embedding(position_ids) embedding token_emb segment_emb position_emb embedding self.layer_norm(embedding) embedding self.dropout(embedding) return embedding这里_position_ids_的生成方式是关键。BERT使用的是可学习的位置编码而不是Transformer原论文里那种基于三角函数的固定编码。意味着position_embedding里的参数会随着训练被更新。在微调阶段如果输入的最大长度超过了预训练时的最大长度比如预训练是512你的任务需要处理600的文本直接报错或者效果崩掉都是有可能的——这也是很多实际场景里把长度限制在128或256的原因之一。4.2 多头自注意力机制这是整个BERT的核心也是让很多人困惑的地方。一句话解释自注意力对于每个词计算它与序列中所有词的关联程度然后按关联程度加权融合所有词的信息。多头注意力的多头本质是不只在单一的表示空间里计算注意力而是把向量拆分到多个子空间并行计算最后拼接。每个头可以关注不同的语义关系比如一个头关注语法依赖另一个头关注指代关系。class MultiHeadAttention(nn.Module): def __init__(self, hidden_size, num_heads, dropout0.1): super().__init__() assert hidden_size % num_heads 0 self.num_heads num_heads self.head_dim hidden_size // num_heads self.hidden_size hidden_size self.q_linear nn.Linear(hidden_size, hidden_size) self.k_linear nn.Linear(hidden_size, hidden_size) self.v_linear nn.Linear(hidden_size, hidden_size) self.out_linear nn.Linear(hidden_size, hidden_size) self.dropout nn.Dropout(dropout) def forward(self, query, key, value, maskNone): batch_size, seq_length, _ query.shape Q self.q_linear(query).view(batch_size, seq_length, self.num_heads, self.head_dim) K self.k_linear(key).view(batch_size, seq_length, self.num_heads, self.head_dim) V self.v_linear(value).view(batch_size, seq_length, self.num_heads, self.head_dim) Q Q.transpose(1, 2) K K.transpose(1, 2) V V.transpose(1, 2) scores Q K.transpose(-2, -1) / (self.head_dim ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attention_weights torch.softmax(scores, dim-1) attention_weights self.dropout(attention_weights) output attention_weights V output output.transpose(1, 2).contiguous().view(batch_size, seq_length, self.hidden_size) output self.out_linear(output) return output这里有个新手比较容易出错的细节scores Q K.transpose(-2, -1) / (self.head_dim ** 0.5)。除以sqrt(head_dim)是标准做法目的是防止点积结果过大导致softmax梯度消失。如果不除以这个缩放因子模型在深层会出现训练不稳定的情况。还有一个容易被忽略的是_contiguous()调用。transpose_操作会改变张量在内存中的存储方式产生一个非连续的内存布局。如果不调用_contiguous()后面的_view()_会直接报错。这类报错信息通常很抽象view size is not compatible with input tensors size很多人被卡在这里很久。4.3 前馈网络与残差连接Transformer块在注意力层后面接了一个两层的前馈网络中间用GELU激活函数。BERT标准配置里这个中间层的大小是768*43072维比输入输出维度大得多。这个设计是有讲究的注意力层负责在词与词之间交互信息而前馈网络层则逐位置地对融合后的信息做非线性变换相当于每个位置的深度思考。class FeedForward(nn.Module): def __init__(self, hidden_size, intermediate_size, dropout0.1): super().__init__() self.linear1 nn.Linear(hidden_size, intermediate_size) self.linear2 nn.Linear(intermediate_size, hidden_size) self.gelu nn.GELU() self.dropout nn.Dropout(dropout) def forward(self, x): x self.linear1(x) x self.gelu(x) x self.dropout(x) x self.linear2(x) return xBERT_pre_training_实现里用了GELU激活函数它和ReLU的区别在于ReLU在x0处有一个硬转折而GELU的曲线是光滑的在小值和负值时不会直接截断为0。理论上这种平滑性让梯度流更顺畅。4.4 编码器层与完整模型组装把上面这些模块组装成一个Transformer块还需要残差连接和层归一化class TransformerEncoderLayer(nn.Module): def __init__(self, hidden_size, num_heads, intermediate_size, dropout0.1): super().__init__() self.attention MultiHeadAttention(hidden_size, num_heads, dropout) self.ffn FeedForward(hidden_size, intermediate_size, dropout) self.norm1 nn.LayerNorm(hidden_size) self.norm2 nn.LayerNorm(hidden_size) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): # 自注意力子层 残差连接 LayerNorm attn_output self.attention(x, x, x, mask) x self.norm1(x self.dropout(attn_output)) # 前馈子层 残差连接 LayerNorm ffn_output self.ffn(x) x self.norm2(x self.dropout(ffn_output)) return x残差连接和LayerNorm在每个子层里缺一不可。残差连接让梯度能直接回传到更浅的层缓解深层网络的梯度消失问题LayerNorm则让每一层的输入分布保持稳定加速收敛。最后把多个Transformer块堆叠起来加上分类头就是完整的BERT模型class BERTForClassification(nn.Module): def __init__(self, vocab_size, hidden_size768, num_layers12, num_heads12, intermediate_size3072, num_labels2, max_length128, dropout0.1): super().__init__() self.embedding BERTEmbedding(vocab_size, hidden_size, max_length, dropout) self.encoder_layers nn.ModuleList([ TransformerEncoderLayer(hidden_size, num_heads, intermediate_size, dropout) for _ in range(num_layers) ]) self.pooler nn.Linear(hidden_size, hidden_size) self.classifier nn.Linear(hidden_size, num_labels) self.dropout nn.Dropout(dropout) def forward(self, input_ids, attention_mask, token_type_ids): embedding_output self.embedding(input_ids, token_type_ids) extended_mask attention_mask.unsqueeze(1).unsqueeze(2) encoder_output embedding_output for layer in self.encoder_layers: encoder_output layer(encoder_output, extended_mask) # 取[CLS]位置的输出作为整个序列的表示 cls_output encoder_output[:, 0, :] pooled_output torch.tanh(self.pooler(cls_output)) pooled_output self.dropout(pooled_output) logits self.classifier(pooled_output) return logits这个架构就是bert-base的完整结构12层Transformer、768维隐藏层、12个注意力头、总共约1.1亿参数。用这个规模微调是必须的。如果显存不够有一个训练技巧冻结前几层Transformer的参数只更新后几层和分类头显存占用会大幅下降微调效果在数据量不大的情况下还说得过去。5. 训练循环优化器选择、学习率调度与loss下降排查模型搭好了接下来是训练环节。BERT微调的套路和训练普通神经网络不太一样主要区别在优化器和学习率策略上。5.1 AdamW优化器和warmup策略BERT官方推荐使用AdamW优化器。它和普通Adam的区别是AdamW把权重衰减weight decay从梯度的滑动平均中分离出来只对参数本身做衰减。在Transformer类模型上这个改动对收敛效果和泛化能力有明显的提升。而warmup策略就更关键了训练初期模型参数是随机初始化的或者预训练权重还没适应新数据如果直接用较大的学习率很容易把参数推到不理想的区域。所以先让学习率从0线性增长到预设值让模型热身再按余弦曲线衰减到0from transformers import get_linear_schedule_with_warmup optimizer torch.optim.AdamW(model.parameters(), lr2e-5, weight_decay0.01) total_steps len(train_dataloader) * num_epochs warmup_steps int(total_steps * 0.1) scheduler get_linear_schedule_with_warmup( optimizer, num_warmup_stepswarmup_steps, num_training_stepstotal_steps )初始学习率的选择上如果加载的是BERT预训练权重做微调经验值是2e-5到5e-5之间。如果是从零训练那需要更大的学习率大概1e-4到5e-4但BERT从零训练的效果通常很差这就是为什么几乎所有实际项目都是加载预训练权重。5.2 完整训练循环模板def train_epoch(model, dataloader, optimizer, scheduler, device): model.train() total_loss 0.0 progress_bar tqdm(dataloader, descTraining) for batch in progress_bar: input_ids batch[input_ids].to(device) attention_mask batch[attention_mask].to(device) token_type_ids batch[token_type_ids].to(device) labels batch[label].to(device) optimizer.zero_grad() logits model(input_ids, attention_mask, token_type_ids) loss nn.CrossEntropyLoss()(logits, labels) loss.backward() # 梯度裁剪防止梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() scheduler.step() total_loss loss.item() progress_bar.set_postfix({loss: f{loss.item():.4f}}) return total_loss / len(dataloader)_grad_clip_这一步很多人会忽略但它在BERT的训练中非常重要。Transformer模型在深层梯度回传时偶尔会出现异常大的梯度不裁剪的话模型参数会瞬间被破坏loss直接变成NaN而且这种损坏是不可逆的——只能重新加载权重再来。5.3 loss不下降的排查链路我实打实遇到过loss卡住不动的情况这里分享完整的排查思路先看数据检查tokenizer产生的input_ids是否全部是同一个值。如果padding策略不对比如max_length设得太小导致所有有效token都被截断只剩[PAD]了模型当然学不到任何东西。打印几条数据人工检查这一步花费不了几分钟但能排除大量低级错误。再看模型输出跑一个batch的前向检查logits的数值范围。如果所有输出都是0附近的值可能是权重初始化问题如果输出极大几十甚至上百可能是LayerNorm写错了位置。验证损失计算算一下初始loss是否接近随机猜测的理论值。二分类的随机猜测loss应该接近ln(2)≈0.693。如果你的初始loss是0.001或者5.0大概率预处理或模型结构有问题。逐层检查梯度backward之后打印每一层梯度的norm。如果某一层的梯度为None或者0说明梯度已经断流要检查残差连接是否正确。最后才怀疑超参学习率太大导致震荡、batch_size太小导致梯度噪声过大这些确实会影响收敛但不该是第一怀疑对象。5.4 验证与模型保存每个epoch之后都要在验证集上评估并且只保留验证集效果最好的模型。评估阶段要用model.eval()和无梯度模式def evaluate(model, dataloader, device): model.eval() predictions [] true_labels [] total_loss 0.0 with torch.no_grad(): for batch in dataloader: input_ids batch[input_ids].to(device) attention_mask batch[attention_mask].to(device) token_type_ids batch[token_type_ids].to(device) labels batch[label].to(device) logits model(input_ids, attention_mask, token_type_ids) loss nn.CrossEntropyLoss()(logits, labels) total_loss loss.item() preds torch.argmax(logits, dim-1) predictions.extend(preds.cpu().tolist()) true_labels.extend(labels.cpu().tolist()) return total_loss / len(dataloader), predictions, true_labels保存模型的时候我习惯把完整的状态字典和配置参数一起存下来torch.save({ model_state_dict: model.state_dict(), config: { vocab_size: tokenizer.vocab_size, hidden_size: 768, num_layers: 12, num_heads: 12, num_labels: 2, max_length: 128 } }, bert_classifier.pt)用torch.save保存整个模型对象虽然方便但会带来兼容性问题——模型类定义的位置变了或者代码重构后load的时候会报错。保存state_dict和config才是最稳妥的方式。6. 加载预训练权重把别人的学习成果接过来前面提到BERT不比从零训练而是加载Google预训练好的权重再做微调。这是整个流程里最体现工程经验的部分。如果模型类是自己写的结构和HuggingFace的_bert-base-chinese_在命名上不一致直接load_state_dict会报键名不匹配的错误。6.1 权重的键名映射HuggingFace的预训练权重键名格式通常是这样的bert.embeddings.word_embeddings.weight、bert.encoder.layer.0.attention.self.query.weight而我上面实现的模型类键名是_embedding.token_embedding.weight_、encoder_layers.0.attention.q_linear.weight。两者的名字完全对不上。解决方法是自己写一个键名映射函数def load_pretrained_weights(model, pretrained_state_dict): mapping { bert.embeddings.word_embeddings.weight: embedding.token_embedding.weight, bert.embeddings.position_embeddings.weight: embedding.position_embedding.weight, bert.embeddings.token_type_embeddings.weight: embedding.segment_embedding.weight, # ... 每个层的映射 } new_state_dict {} for k, v in pretrained_state_dict.items(): mapped_key mapping.get(k) if mapped_key is not None: new_state_dict[mapped_key] v # 加载missing_keys里是分类头等随机初始化的参数 model.load_state_dict(new_state_dict, strictFalse)这个方案可行但从维护成本和错误率的角度看我后来都直接用HuggingFace封装好的_bert-base-chinese_模型来做骨干然后自己在上面套输出头。这样就告别了繁琐的键名映射代码也更稳。自己实现模型结构最大的价值在于理解原理真实做项目时我推荐的方法是用自己的代码做验证理解用HuggingFace做实际工程。6.2 两个典型的权重加载报错常见的报错之一是_missing keys_表示模型里有预训练权重里没有的键通常是新加的分类头参数属于正常现象。另一种是_unexpected keys_表示预训练权重里有的键在模型里找不到通常是键名映射没写全。还有一个非常隐蔽的坑词表大小不一致。如果你的tokenizer还是默认的_bert-base-chinese_那词表是21128个token预训练权重也是按照这个尺寸生成的两边匹配没问题。但如果你中途往tokenizer里加过自定义token比如把[MASK]之外又加了 之类的符号词表大小变成了21130嵌入层的维度就错位了load_state_dict直接报size mismatch。这个问题可以让它的报错信息搞得你怀疑人生排查半天才发现是多加了特殊token。7. 推理部署训练结束才是真正的开始模型训好之后部署阶段也有不少坑。这里分享一个完整的单条文本预测函数def predict(model, tokenizer, text, device, max_length128): model.eval() encoding tokenizer( text, max_lengthmax_length, paddingmax_length, truncationTrue, return_tensorspt ) input_ids encoding[input_ids].to(device) attention_mask encoding[attention_mask].to(device) token_type_ids encoding[token_type_ids].to(device) with torch.no_grad(): logits model(input_ids, attention_mask, token_type_ids) probs torch.softmax(logits, dim-1) pred torch.argmax(logits, dim-1).item() return pred, probs.squeeze().cpu().tolist()这个函数看起来简单但有几处值得注意。模型要切换到eval模式这个动作不只是语义上的告诉模型别训练了而是真的会改变BatchNorm和Dropout的行为。eval模式下Dropout会被关闭。如果不调用model.eval()同样的文本在不同时刻跑出来结果会有细微波动这在线上推理时是致命的。用no_grad包裹推理逻辑原因很简单——省内存。在torch.no_grad()上下文中框架不会构建计算图也不需要保存中间激活值长文本推理的显存占用可能只有普通方式的一半。如果需要在CPU上做推理可以用torch.compile加上量化操作来加速。BERT这类Transformer模型在CPU上的推理速度向来是个痛点加速方向主要是两个一是转为ONNX格式用ONNX Runtime推理二是用动态量化把float32权重变成int8。两种方法我都试过ONNX在长文本上的提速大约2到3倍量化则可能带来一些精度损失需要看业务场景接受程度。8. 数据增强与效果提升的进阶思路如果微调后模型效果还是不够理想有几个方向值得尝试按性价比从高到低排列。8.1 用领域语料做增量预训练通用BERT对专业领域的文本理解有限比如医疗、法律、金融这些领域有大量的专业术语和特有表达。常规做法是加载预训练权重之后用领域内的大规模无标注文本继续训练MLM任务然后再用标注数据做监督微调。这个技术在HuggingFace里封装得很好只需要用_DataCollatorForLanguageModeling_去构造训练样本。实战中即使只用10万条无标注的领域文本做增量预训练在特定任务的准确率上也能看到可见的提升。8.2 集成多模型投票如果单模型的准确率卡在瓶颈无法突破第三板斧是模型集成。把BERT、RoBERTa、ERNIE等不同预训练模型微调出来在推理时对预测概率取平均或者做投票。这个方法近乎无脑但效果稳定——每个模型在结构上有差异错误的相关性相对较低投票后泛化能力会有所提升。代价是推理耗时成倍增加部署成本也水涨船高。8.3 判别性学习率还有一个容易出效果的技巧是对模型不同层设置不同的学习率。预训练层已经有很好的通用语义理解微调时不需要大动干戈而随机初始化的分类头需要更大的更新力度。实现方式很简单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) and classifier not in n], weight_decay: 0.01, lr: 2e-5 }, { params: [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay) and classifier not in n], weight_decay: 0.0, lr: 2e-5 }, { params: [p for n, p in model.named_parameters() if classifier in n], weight_decay: 0.01, lr: 1e-4 } ]这个做法的合理性在于不同的权重衰减和更新速度让模型在不同层之间有一个学习速率差预训练层学得保守一点输出头学得激进一点。实际操作里效果提升虽然称不上惊艳但往往能在验证集上带来零点几个百分点的准确率提升。9. 高频踩坑清单这些错误我几乎都犯过整理几个最典型的错误按出现频率排序。第一DataLoader的batch维度处理。HuggingFace的tokenizer默认返回形状为[1, seq_len]的张量如果你在Dataset里忘了squeezeDataLoader组合后变成[batch_size, 1, seq_len]模型forward时维度对不上报错信息还特别难懂。建议在第一次写Dataset类的时候就在__getitem__里squeeze干净。第二multi-GPU训练时的device错误。用DataParallel包装模型后DDP会自动把输入分发到各个GPU上但如果你手动把数据.to(device)这里的device往往只是主GPU的编号模型在其它GPU上处理的数据反而对不上。使用_DataParallel时不要手动去.to(device)直接传CPU上的数据即可。第三中文数据里混入了特殊字符。比如全角空格、零宽字符、不可见字符这些在文本清理阶段就要处理掉。因为词表里没有这些字符分词器会自动把它们映射成[UNK]如果你的数据里[UNK]比例偏高模型语义理解能力会明显变差。第四显存OOM后没有及时清空CUDA缓存。尤其是训练过程中不断有点测或可视化操作显存碎片会越积越多。在训练循环的适当位置插入_torch.cuda.empty_cache()_让显存碎片能被回收。第五用了不相容的测试数据分布。模型在验证集上效果不错但上线后预测结果一塌糊涂最常见的原因是训练数据分布和真实业务数据分布不一致。比如训练数据都是经过清理的简洁文本而线上数据里充满了广告文本、特殊符号、无关信息。解决思路是采集一部分真实线上数据做增量微调而不是只依赖人工清洗过的那部分。本文还有配套的精品资源点击获取
RELATED — 相关阅读

相关资讯

LATEST — 最新资讯

最新发布

TODAY — 本日精选

新闻

WEEKLY — 本周精选

新闻

MONTHLY — 本月精选

新闻