1. 为什么要在 Keras 生态里死磕 ALBERT做 NLP 工程落地的朋友大概率都经历过这个纠结BERT 效果好但太沉RoBERTa 精度高但显存吃紧蒸馏出来的小模型又经常在业务数据上掉点。ALBERT 刚出来那阵子我其实没太在意觉得无非又是一个参数少一点的变体直到有一次在一个文本分类项目上被显存卡死——单卡 11GBERT-base 开到 batch size 16 就 OOM换 ALBERT 之后同样的卡能跑到 batch size 48推理延迟直接砍掉将近一半精度只掉了不到 0.5 个点。从那之后我就把 ALBERT 纳入了常规武器库。这篇东西主要聊的是怎么用keras-bert这个库去调用 ALBERT把文本分类、文本多标签分类、序列标注这三类最常见的任务一次性打通。keras-bert是 bojone 维护的一个基于 Keras 的 BERT 系列模型加载与微调工具它的好处是 API 极其简洁加载预训练权重、接下游任务、导出模型这几步都很顺不像 HuggingFace 那套需要理解一堆 Trainer、DataCollator 的概念。对于习惯 Keras 那套fit/predict流程的人来说上手成本几乎为零。适合谁看如果你已经跑通过 BERT 的基本微调想找一个更省显存的替代方案或者你手头有 Keras/TensorFlow 的工程栈不想为了一个分类任务再引入 PyTorch 依赖再或者你需要同时处理单标签、多标签、序列标注三种任务想找一套统一的代码骨架——那这篇基本能覆盖你的需求。下面我会把设计思路、核心细节、完整实操、踩坑记录都摊开讲代码可以直接抄。2. ALBERT 与 keras-bert 的组合逻辑拆解2.1 ALBERT 到底省在哪参数分解与跨层共享很多人以为 ALBERT 是小号 BERT其实它的核心改动就两处但这两处直接决定了它的显存和速度优势。第一处是词嵌入参数分解Factorized Embedding Parameterization。标准 BERT 的词嵌入维度 E 和隐藏层维度 H 是相等的base 版都是 768。这意味着词表大小 V 乘以 768 就是嵌入矩阵的参数量V 通常三万左右光嵌入层就两千多万参数。ALBERT 把 E 和 H 解耦让 E 远小于 H比如 E128、H768嵌入矩阵先投影到 128 维再通过一个 128×768 的矩阵升到隐藏维度。这样嵌入层参数量从 V×H 降到 V×E E×H直接砍掉一大截。第二处是跨层参数共享Cross-layer Parameter Sharing。BERT-base 有 12 层 Transformer每层都有独立的注意力权重和前馈网络权重。ALBERT 让这 12 层共享同一套参数等于把 12 层的参数量压缩成 1 层。这一招对显存的节省是实打实的因为反向传播时只需要维护一份参数的梯度。注意参数共享带来一个副作用——ALBERT 的训练收敛比 BERT 慢一些因为每层都在更新同一套权重梯度信号会互相干扰。实践中我一般会把学习率调得比 BERT 微调时略小warmup 步数拉长一点。2.2 为什么选 keras-bert 而不是别的加载方式市面上加载 ALBERT 的路子主要有三条HuggingFace Transformers、Google 官方 TF-Hub、以及 keras-bert。我选 keras-bert 的理由很实际依赖干净。它只依赖 Keras 和 TensorFlow不需要额外装 tokenizers、datasets 那一堆。对于已经用 TF 做生产的团队少一个依赖就少一份部署麻烦。API 直白。load_vocab读词表TokenEmbedding加载权重build_transformer_model搭骨架model.fit微调四步走完。没有 config 对象层层嵌套。下游任务封装好。keras_bert.layers里有现成的Masked、Extract这些层序列标注任务里取每个 token 的隐藏状态特别方便。当然它也有短板预训练权重的下载需要自己准备不像 HF 那样from_pretrained一个字符串就搞定。但 ALBERT 的权重文件本来就不大base 版压缩后才十几兆手动下载一次的事。2.3 三类任务的统一建模思路文本分类、多标签分类、序列标注表面看是三件事但在 BERT 类模型的微调框架下它们的差异只在输出层和损失函数任务类型取哪个位置的输出输出层损失函数激活单标签分类[CLS] 位置Dense(num_classes)categorical_crossentropysoftmax多标签分类[CLS] 位置Dense(num_labels)binary_crossentropysigmoid序列标注每个 token 位置Dense(num_tags)sparse_categorical_crossentropysoftmax逐位置所以整套代码可以共用同一个骨架只在最后接不同的 head。这也是我下面实操部分组织代码的主线——先搭一个通用的 ALBERT 加载器再分别接三种 head。3. 环境准备与 ALBERT 权重加载实操3.1 依赖安装与版本对齐先把环境理清楚。keras-bert 对版本比较敏感我实测下来这套组合最稳pip install tensorflow2.6.0 pip install keras2.6.0 pip install keras-bert0.89.0 pip install numpy pandas scikit-learn提示keras-bert 0.89.0 之后的版本对 TF 2.x 的兼容性有波动如果你用的是 TF 2.10 以上可能会遇到tf.keras和独立keras包冲突的问题。稳妥做法是锁定 TF 2.6 到 2.8 这个区间。如果你非要用更新的 TF有个绕法是在导入前设置环境变量TF_USE_LEGACY_KERAS1让 TF 走老的 Keras 实现。这个我踩过坑不设的话build_transformer_model会报找不到某些层的错。3.2 ALBERT 权重与词表的准备ALBERT 的预训练权重和配置文件需要提前下好。以 base 版为例你需要这几个文件albert_config.json模型结构配置albert_model.ckpt权重文件注意是 TF 格式的 checkpoint不是 PyTorch 的 binvocab.txt词表目录结构建议这样放pretrained/ albert_base/ albert_config.json albert_model.ckpt.index albert_model.ckpt.data-00000-of-00001 vocab.txt加载的代码就三行import os from keras_bert import load_vocab, TokenEmbedding, build_transformer_model config_path pretrained/albert_base/albert_config.json checkpoint_path pretrained/albert_base/albert_model.ckpt vocab_path pretrained/albert_base/vocab.txt token_dict load_vocab(vocab_path)这里有个细节load_vocab返回的是一个有序字典key 是 tokenvalue 是索引。后面做 tokenize 的时候要用到它。3.3 构建 ALBERT 骨架的关键参数build_transformer_model是核心函数参数不少我挑几个关键的讲from keras_bert import build_transformer_model model build_transformer_model( config_pathconfig_path, checkpoint_pathcheckpoint_path, modelalbert, applicationalbert, return_keras_modelFalse, )modelalbert告诉库用 ALBERT 的结构而不是 BERT。这个参数不写的话默认按 BERT 解析会报层名对不上的错。applicationalbert指定下游应用的接口风格。ALBERT 的 [CLS] 输出和 BERT 略有差异这个参数会影响Extract层的取法。return_keras_modelFalse返回的是 Keras-BERT 的封装对象方便后面调.model拿底层 Keras 模型。如果你直接要 Keras Model设成 True 也行。搭好之后model.model就是那个可以接 head 的 Keras 模型输入是两个token ids 和 segment ids。4. 文本分类任务的完整实现4.1 数据预处理与 tokenize文本分类是最简单的场景我拿一个二分类的示例数据来演示。假设数据长这样data [ (这个产品的质量非常好用了一个月没出问题, 1), (收到货就坏了客服还不理人, 0), # ... 更多样本 ]tokenize 用 keras-bert 自带的Tokenizerfrom keras_bert import Tokenizer tokenizer Tokenizer(token_dict) def encode(text, max_len128): tokens tokenizer.tokenize(text) # 截断到 max_len - 2留出 [CLS] 和 [SEP] tokens tokens[:max_len - 2] # 首尾加特殊符号 tokens [[CLS]] tokens [[SEP]] # 补齐到 max_len token_ids tokenizer.tokens_to_ids(tokens) segment_ids [0] * len(token_ids) # padding token_ids [0] * (max_len - len(token_ids)) segment_ids [0] * (max_len - len(segment_ids)) return token_ids, segment_ids注意tokenizer.tokenize内部已经做了 basic tokenize 和 wordpiece 切分中文是按字切的英文按 subword 切。如果你的文本里有大量英文专业术语建议先做一轮小写化因为 ALBERT 的词表里英文 token 基本都是小写的。4.2 接分类 head 并编译取 [CLS] 位置的输出接一个 Dense 层from tensorflow.keras.layers import Dense, Dropout from tensorflow.keras.models import Model # 取 [CLS] 对应的隐藏状态 cls_output model.model.output[0][:, 0, :] # 或者用 keras_bert 的 Extract 层 # from keras_bert import Extract # cls_output Extract()(model.model.output) x Dropout(0.3)(cls_output) output Dense(2, activationsoftmax)(x) final_model Model(model.model.input, output) final_model.compile( optimizertf.keras.optimizers.Adam(learning_rate2e-5), losscategorical_crossentropy, metrics[accuracy] )这里model.model.output[0]是序列输出形状是(batch, seq_len, hidden)取[:, 0, :]就是 [CLS] 位置。model.model.output[1]是 pooled output但 ALBERT 的 pooled 输出有时候不太靠谱我一般直接用 [CLS] 的隐藏状态。学习率设 2e-5 是微调的常规操作ALBERT 因为参数共享我试过 1e-5 到 3e-5 之间2e-5 在多数数据集上比较稳。4.3 训练与评估的实操细节训练前把标签转成 one-hotimport numpy as np from tensorflow.keras.utils import to_categorical X1, X2, Y [], [], [] for text, label in data: t1, t2 encode(text) X1.append(t1) X2.append(t2) Y.append(label) X1 np.array(X1) X2 np.array(X2) Y to_categorical(np.array(Y), num_classes2) final_model.fit( [X1, X2], Y, batch_size32, epochs5, validation_split0.1 )batch size 32 在 11G 显存上跑 ALBERT-base 绰绰有余我甚至试过 64 也没爆。epochs 一般 3 到 5 就够再多容易过拟合尤其是数据量小于一万条的时候。实操心得ALBERT 在小数据集上过拟合比 BERT 更快因为参数共享导致模型容量虽然小了但梯度更新更集中。我的做法是加一个早停回调监控验证集 losspatience 设 2。5. 文本多标签分类的改造要点5.1 多标签与单标签的本质差异多标签分类的场景很常见比如一条新闻同时属于科技和财经两个类目或者一段用户反馈同时涉及物流和售后两个问题。它和单标签的核心区别在于标签之间不互斥所以不能用 softmax要用 sigmoid 对每个标签独立判断。损失函数从categorical_crossentropy换成binary_crossentropy输出层神经元数量等于标签总数每个神经元输出一个 0 到 1 之间的概率。5.2 标签编码与输出层调整假设有 6 个标签一条样本的标签是[1, 0, 1, 0, 0, 1]表示同时命中第 1、3、6 个类目。num_labels 6 x Dropout(0.3)(cls_output) output Dense(num_labels, activationsigmoid)(x) final_model Model(model.model.input, output) final_model.compile( optimizertf.keras.optimizers.Adam(learning_rate2e-5), lossbinary_crossentropy, metrics[accuracy] )标签直接构造成 multi-hot 向量就行不需要 one-hotY np.array([ [1, 0, 1, 0, 0, 1], [0, 1, 0, 0, 1, 0], # ... ])5.3 阈值选择与评估指标多标签分类有个绕不开的问题预测阈值怎么定。默认 0.5 在很多场景下不是最优的。我的做法是在验证集上扫一遍阈值从 0.1 到 0.9看哪个阈值下 F1 最高。from sklearn.metrics import f1_score preds final_model.predict([X_val1, X_val2]) best_thresh, best_f1 0.5, 0 for thresh in np.arange(0.1, 0.9, 0.05): binary_preds (preds thresh).astype(int) f1 f1_score(Y_val, binary_preds, averagemicro) if f1 best_f1: best_f1 f1 best_thresh thresh print(f最佳阈值: {best_thresh:.2f}, 对应 F1: {best_f1:.4f})注意多标签任务里accuracy这个指标参考价值很低因为大部分标签都是 0模型全预测 0 也能拿到很高的 accuracy。一定要看 micro-F1 或 macro-F1。6. 序列标注任务的实现与对齐问题6.1 序列标注的输出层设计序列标注NER、POS 等和分类最大的不同是每个 token 都要出一个预测而不是整句出一个。所以输出层要作用在序列维度上from tensorflow.keras.layers import TimeDistributed num_tags 9 # 比如 BIO 标注下的实体类型数 sequence_output model.model.output[0] # (batch, seq_len, hidden) x Dropout(0.3)(sequence_output) output Dense(num_tags, activationsoftmax)(x) final_model Model(model.model.input, output) final_model.compile( optimizertf.keras.optimizers.Adam(learning_rate2e-5), losssparse_categorical_crossentropy, metrics[accuracy] )这里用sparse_categorical_crossentropy是因为标签是整数索引不需要 one-hot。如果你用 one-hot 标签就换成categorical_crossentropy。6.2 标签对齐最容易出错的地方序列标注的坑几乎全在标签对齐上。BERT 类模型会把一个词切成多个 subword但你的标注通常是在词级别。比如北京市如果被切成北京和市两个 token那市这个 token 的标签怎么给常见做法有三种只在第一个 subword 上计算损失其余 subword 的标签设为 -1忽略。这是最稳妥的也是我推荐的。给所有 subword 都赋同样的标签。简单但会引入噪声。用[CLS]和[SEP]之外的 token 做平均。复杂且收益不明显。我一般用第一种配合sample_weight或者直接把忽略位置的标签设成一个特殊值在 loss 里 mask 掉。def align_labels(tokens, labels, max_len128): tokens 是 tokenize 后的列表labels 是词级别的标签列表 aligned [] for token, label in zip(tokens, labels): # 第一个 subword 给真实标签后续 subword 给 -1 aligned.append(label) for _ in range(len(token) - 1): aligned.append(-1) # 加 [CLS] 和 [SEP] 的标签 aligned [-1] aligned[:max_len - 2] [-1] aligned [-1] * (max_len - len(aligned)) return aligned然后在 loss 里把 -1 的位置忽略掉。Keras 的sparse_categorical_crossentropy默认不忽略 -1需要自己写一个 maskdef masked_loss(y_true, y_pred): mask tf.cast(tf.not_equal(y_true, -1), tf.float32) y_true tf.where(tf.equal(y_true, -1), tf.zeros_like(y_true), y_true) loss tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred) loss loss * mask return tf.reduce_sum(loss) / tf.reduce_sum(mask)6.3 解码与后处理预测出来的是每个 token 的标签概率需要转成实体。以 BIO 标注为例def decode_tags(tokens, tag_probs, id2tag): tags [id2tag[i] for i in np.argmax(tag_probs, axis-1)] entities [] current None for token, tag in zip(tokens, tags): if tag.startswith(B-): if current: entities.append(current) current {type: tag[2:], tokens: [token]} elif tag.startswith(I-) and current and current[type] tag[2:]: current[tokens].append(token) else: if current: entities.append(current) current None if current: entities.append(current) return entities实操心得序列标注任务里[CLS]和[SEP]位置的预测一定要忽略否则解码时会冒出莫名其妙的实体。另外 padding 位置的预测也要忽略不然短句后面会跟一堆垃圾标签。7. 常见问题与排查速查7.1 加载权重报错层名不匹配最常见的报错是NotFoundError: Key bert/embeddings/... not found in checkpoint。这通常是因为model参数没设成albert库按 BERT 的层名去找权重自然找不到。检查build_transformer_model的model和application两个参数是否都设对了。另一个可能是权重文件本身是 PyTorch 格式的.binkeras-bert 只认 TF 的 checkpoint。需要先用转换脚本转一下或者直接下 TF 版本的权重。7.2 显存够但训练极慢ALBERT 虽然参数少但计算量并没有等比下降因为前向传播还是要过 12 层。如果训练慢先检查是不是在 CPU 上跑。另外tf.data的 pipeline 如果没开prefetch数据加载会成为瓶颈。dataset tf.data.Dataset.from_tensor_slices(([X1, X2], Y)) dataset dataset.batch(32).prefetch(tf.data.AUTOTUNE)7.3 多标签任务 F1 上不去先别急着调模型检查三件事标签分布是否极度不均衡、阈值是否合适、正样本是否太少。如果某个标签只有几十个正样本模型基本学不动这时候要么做数据增强要么用 focal loss 替代 binary crossentropy。7.4 序列标注预测全是 O 标签这是典型的类别不均衡问题。NER 数据里 O 标签通常占 80% 以上模型偷懒全预测 O 就能拿到很高的 accuracy。解决办法是给非 O 标签加权或者在 loss 里对 O 标签降权。class_weights { 0: 0.2, # O 标签降权 1: 2.0, # B-XXX 加权 2: 2.0, # I-XXX 加权 # ... } final_model.fit(..., class_weightclass_weights)7.5 问题速查表现象可能原因排查方向加载权重报 Key 不存在model 参数没设 albert检查 build_transformer_model 参数训练 loss 不下降学习率过大或过小试 1e-5 到 3e-5验证集 loss 先降后升过拟合加 Dropout、早停、减 epoch多标签 F1 低阈值不合适扫阈值找最优序列标注全 O类别不均衡加 class_weight推理速度慢没开 GPU 或 batch 太小检查 device、增大 batch8. 我踩过的几个真实坑第一个坑是词表大小和 config 对不上。有次我拿了一个自己扩充过词表的 ALBERT 权重但 config 里的vocab_size没改结果加载时嵌入层形状对不上报了个很隐晦的 shape mismatch。后来养成习惯加载前先核对 config 里的vocab_size和vocab.txt的行数是否一致。第二个坑是segment_ids 全填 0。单句任务里 segment_ids 确实全是 0但如果你做的是句子对任务比如文本蕴含第二句的 segment_ids 要填 1。我一开始偷懒全填 0模型在句子对任务上效果差了一大截排查了半天才发现。第三个坑是max_len 设得太大。ALBERT 支持到 512但大部分分类任务 128 就够了。我有个项目图省事设了 512结果显存占用翻倍训练速度慢了三倍精度还没提升。后来改成动态 padding按 batch 内最长句补齐效率提升明显。第四个坑是忘了 freeze 底层。小数据集上微调 ALBERT如果所有层都放开训练很容易过拟合。我的做法是前 6 层 freeze只训后 6 层和 head效果反而更稳。for layer in model.model.layers[:6]: layer.trainable False这套代码骨架我在好几个项目上复用过了从情感分类到意图识别再到实体抽取改的基本就是 head 和 loss。ALBERT 加 keras-bert 这个组合对于显存紧张又不想牺牲太多精度的场景确实是个很务实的选择。