1. 为什么要在 Keras 生态里死磕 ALBERT

做 NLP 工程落地的朋友大概率都经历过这个纠结：BERT 效果好但太沉，RoBERTa 精度高但显存吃紧，蒸馏出来的小模型又经常在业务数据上掉点。ALBERT 刚出来那阵子我其实没太在意，觉得无非又是一个"参数少一点"的变体，直到有一次在一个文本分类项目上被显存卡死——单卡 11G，BERT-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，比如 E=128、H=768，嵌入矩阵先投影到 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 类模型的微调框架下，它们的差异只在输出层和损失函数：

所以整套代码可以共用同一个骨架，只在最后接不同的 head。这也是我下面实操部分组织代码的主线——先搭一个通用的 ALBERT 加载器，再分别接三种 head。

3. 环境准备与 ALBERT 权重加载实操

3.1 依赖安装与版本对齐

先把环境理清楚。keras-bert 对版本比较敏感，我实测下来这套组合最稳：

pip install tensorflow==2.6.0 pip install keras==2.6.0 pip install keras-bert==0.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_KERAS=1，让 TF 走老的 Keras 实现。这个我踩过坑，不设的话build_transformer_model会报找不到某些层的错。

3.2 ALBERT 权重与词表的准备

ALBERT 的预训练权重和配置文件需要提前下好。以 base 版为例，你需要这几个文件：

• albert_config.json：模型结构配置
• albert_model.ckpt：权重文件（注意是 TF 格式的 checkpoint，不是 PyTorch 的 bin）
• vocab.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 是 token，value 是索引。后面做 tokenize 的时候要用到它。

3.3 构建 ALBERT 骨架的关键参数

build_transformer_model是核心函数，参数不少，我挑几个关键的讲：

from keras_bert import build_transformer_model model = build_transformer_model( config_path=config_path, checkpoint_path=checkpoint_path, model='albert', application='albert', return_keras_model=False, )

• model='albert'：告诉库用 ALBERT 的结构，而不是 BERT。这个参数不写的话默认按 BERT 解析，会报层名对不上的错。
• application='albert'：指定下游应用的接口风格。ALBERT 的 [CLS] 输出和 BERT 略有差异，这个参数会影响Extract层的取法。
• return_keras_model=False：返回的是 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 自带的Tokenizer：

from keras_bert import Tokenizer tokenizer = Tokenizer(token_dict) def encode(text, max_len=128): 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, activation='softmax')(x) final_model = Model(model.model.input, output) final_model.compile( optimizer=tf.keras.optimizers.Adam(learning_rate=2e-5), loss='categorical_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-hot：

import 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_classes=2) final_model.fit( [X1, X2], Y, batch_size=32, epochs=5, validation_split=0.1 )

batch size 32 在 11G 显存上跑 ALBERT-base 绰绰有余，我甚至试过 64 也没爆。epochs 一般 3 到 5 就够，再多容易过拟合，尤其是数据量小于一万条的时候。

实操心得：ALBERT 在小数据集上过拟合比 BERT 更快，因为参数共享导致模型容量虽然小了，但梯度更新更集中。我的做法是加一个早停回调，监控验证集 loss，patience 设 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, activation='sigmoid')(x) final_model = Model(model.model.input, output) final_model.compile( optimizer=tf.keras.optimizers.Adam(learning_rate=2e-5), loss='binary_crossentropy', metrics=['accuracy'] )

标签直接构造成 multi-hot 向量就行，不需要 one-hot：

Y = 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, average='micro') 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, activation='softmax')(x) final_model = Model(model.model.input, output) final_model.compile( optimizer=tf.keras.optimizers.Adam(learning_rate=2e-5), loss='sparse_categorical_crossentropy', metrics=['accuracy'] )

这里用sparse_categorical_crossentropy是因为标签是整数索引，不需要 one-hot。如果你用 one-hot 标签，就换成categorical_crossentropy。

6.2 标签对齐：最容易出错的地方

序列标注的坑几乎全在标签对齐上。BERT 类模型会把一个词切成多个 subword，但你的标注通常是在词级别。比如"北京市"如果被切成"北京"和"市"两个 token，那"市"这个 token 的标签怎么给？

常见做法有三种：

1. 只在第一个 subword 上计算损失，其余 subword 的标签设为 -1（忽略）。这是最稳妥的，也是我推荐的。
2. 给所有 subword 都赋同样的标签。简单但会引入噪声。
3. 用[CLS]和[SEP]之外的 token 做平均。复杂且收益不明显。

我一般用第一种，配合sample_weight或者直接把忽略位置的标签设成一个特殊值，在 loss 里 mask 掉。

def align_labels(tokens, labels, max_len=128): """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，需要自己写一个 mask：

def 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 格式的.bin，keras-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_weight=class_weights)

7.5 问题速查表

8. 我踩过的几个真实坑

第一个坑是词表大小和 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 这个组合，对于显存紧张又不想牺牲太多精度的场景，确实是个很务实的选择。