兼容 Tensorflow1.x/2.x 的高层封装 (Transformer/GPT-2/BERT/ALBERT/UniLM/XLNet/ELECTRA 等)，使用简单的代码完成语言模型/文本分类/文本生成/命名实体识别/机器阅读理解/机器翻译/序列标注/知识蒸馏任务。适用于 NLP 从业者。
- 高效调用：三行代码完成训练及推理
- 高效运行：一行代码设置多进程/多 GPU 并行
- 品类丰富：支持 40+ 模型类
- 高分保证：提供对比学习、对抗式训练等多项训练技巧
- 可供部署：导出模型 PB 文件，供线上部署
git clone https://github.com/geyingli/unif
cd unif
python3 setup.py install --user若需卸载，通过 pip3 uninstall uf 即可。
import uf
# 建模
model = uf.BERTClassifier(config_file="./ref/bert_config.json", vocab_file="./ref/vocab.txt")
# 定义训练样本
X, y = ["久旱逢甘露", "他乡遇故知"], [1, 0]
# 训练
model.fit(X, y)
# 推理
print(model.predict(X))| 领域 | API | 说明 |
|---|---|---|
| 语言模型 | BERTLM |
结合 MLM 和 NSP 任务，随机采样自下文及其他文档 |
RoBERTaLM |
仅 MLM 任务，采样至文档结束 | |
ALBERTLM |
结合 MLM 和 SOP，随机采样自上下文及其他文档 | |
ELECTRALM |
结合 MLM 和 RTD，生成器与判别器联合训练 | |
VAELM |
可生成语言文本负样本，也可提取向量用于聚类 | |
GPT2LM |
自回归式文本生成 | |
UniLM |
结合双向、单向及 Seq2Seq 建模的全能语言模型 | |
UniLMPrompt |
加入 prompt，进一步实现语言模型与下游任务的统一 | |
| 文本分类 / 单label | TextCNNClassifier |
小而快 |
RNNClassifier |
经典 RNN/LSTM/GRU | |
BiRNNClassifier |
双向获取更优表征 | |
BERTClassifier |
- | |
XLNetClassifier |
- | |
ALBERTClassifier |
- | |
ELECTRAClassifier |
- | |
WideDeepClassifier |
通过 Wide & Deep 架构融合句子级别特征 | |
SemBERTClassifier |
通过 SemBERT 架构融合字级别的特征 | |
UDAClassifier |
结合一致性学习的半监督学习算法 | |
PerformerClassifier |
基于正交随机向量快速计算 attention，实现加速 | |
| 文本分类 / 多label | BERTBinaryClassifier |
- |
XLNetBinaryClassifier |
- | |
ALBERTBinaryClassifier |
- | |
ELECTRABinaryClassifier |
- | |
| 回归 | BERTRegressor |
- |
WideDeepRegressor |
通过 Wide & Deep 架构融合句子级别特征 | |
| 序列标注 | BERTSeqClassifier |
- |
ALBERTSeqClassifier |
- | |
ELECTRASeqClassifier |
- | |
BERTSeqCrossClassifier |
序列标注与文本分类相结合的多任务学习 | |
| 命名实体识别 | BERTNER |
- |
BERTCRFNER |
结合 CRF | |
BERTCRFCascadeNER |
实体识别与分类同时进行的级联架构 | |
| 机器阅读理解 | BERTMRC |
- |
ALBERTMRC |
- | |
SANetMRC |
引入 sentence attention | |
BERTVerifierMRC |
抽取 answer span 的同时判断可答性 | |
RetroReaderMRC |
抽取 answer span 的同时判断可答性 | |
| 机器翻译 | TransformerMT |
共享词表，标准 Seq2Seq 架构 |
| 模型蒸馏 | TinyBERTClassifier |
大幅压缩模型参数，提速十倍以上 |
TinyBERTBinaryClassifier |
- | |
FastBERTClassifier |
动态推理，易分样本提前离开模型 | |
| 图像分类 / 单label | PNasNetClassifier |
基于 AutoML 搜索最佳网络结构 |
实际上建模所需的参数不在少数，因模型而异。为了简便起见，大多数设置了默认值。了解每一项参数的含义是十分有必要的。参数的命名与原论文保持一致，如果遇到不明白的参数，除了看源代码外，可以前往原论文寻找答案。以 BERTClassifier 为例，包含以下参数：
model = uf.BERTClassifier(
config_file, # json格式的配置文件，通常可以在预训练参数包里找到
vocab_file, # 一行一个字/词的txt文件
max_seq_length=128, # 切词后的最大序列长度
label_size=2, # label取值数
init_checkpoint=None, # 预训练参数的路径或目录
output_dir="./output", # 输出文件导出目录
gpu_ids="0,1,3,5", # GPU代号 (为空代表不使用GPU; 如果使用的是Nvidia显卡，需要预先安装CUDA及cuDNN，而后可以通过`nvidia-smi`指令查看可用GPU代号)
drop_pooler=False, # 建模时是否跳过 pooler 层
do_lower_case=True, # 英文是否进行小写处理
truncate_method="LIFO", # 输入超出`max_seq_length`时的截断方式 (LIFO:尾词先弃, FIFO:首词先弃, longer-FO:长文本先弃)
)模型使用完毕后，若需要清理内存，可以使用 del model 删除模型，或通过 model.reset() 对模型进行重置。
同样，训练也包含一些可自行调节的参数，有些参数甚至十分关键：
model.fit(
X=X, # 输入列表
y=y, # 输出列表
sample_weight=None, # 样本权重列表，放空则默认每条样本权重为1.0
X_tokenized=None, # 输入列表 (已预先分词处理的`X`)
batch_size=32, # 每训练一步使用多少数据
learning_rate=5e-05, # 学习率
target_steps=None, # 放空代表直接不间断地训练到`total_steps`；否则为训练停止的位置
total_steps=-3, # 模型计划训练的总步长，决定了学习率的变化曲线；正数，如1000000，代表训练一百万步；负数，如-3，代表根据数据量循环三轮的总步长
warmup_ratio=0.1, # 训练初期学习率从零开始，线性增长到`learning_rate`的步长范围；0.1代表在前10%的步数里逐渐增长
print_per_secs=1, # 多少秒打印一次训练信息
save_per_steps=1000, # 多少步保存一次模型参数
**kwargs, # 其他训练相关参数，如分层学习率等，下文将介绍
)训练过程中，通常需要设立多个断点进行模型验证，决定是否停止训练。target_steps 正是为设置断点而存在的。以下是使用示例：
num_loops = 10 # 假设训练途中一共设置10个断点
num_epochs = 6 # 假设总共训练6轮
for loop_id in range(10):
model.fit(
X, y,
target_steps=-((loop_id + 1) * num_epochs / num_loops), # 训练断点 (-0.6, -1.2, ...)
total_steps=-num_epochs, # 训练全长 (-6)
)
print(model.score(X_dev, y_dev)) # 验证模型
model.localize(f"breakpoint.{loop_id}", into_file=".unif") # 保存模型配置到`into_file` (同时保存模型参数到`output_dir`)多次验证后表现最佳的断点，可以通过 restore 函数取用：
model = uf.restore("breakpoint.7", from_file=".unif")从以上代码不难看出，localize 和 restore 函数是模型管理的利器。
fit 函数内部包含了两个步骤：
-
对输入进行预处理，转换为模型可接受的输入 (e.g. 整数/浮点数矩阵)
-
训练模型
当数据量变得庞大时，例如百万级，数据预处理可能要消耗十几二十分钟，这段期间 GPU 处于闲置状态，无疑是对资源的浪费，可以通过开启多进程处理加速这一过程 (注：不会对第二步模型训练加速)：
with uf.MultiProcess():
X, y = ... # 读取数据
model.fit(X, y)由于 python 中存在 PIL锁，每一个进程只能使用一个 CPU，那么多进程唤醒其他 CPU 的本质是对当前进程进行复制。因此需要注意的是，最好在大批量数据读到程序内存以前开启 MultiProcess，而不要在之后，否则每一个复制的进程都会拷贝一份完整数据，造成不必要的内存占用。
当数据规模进一步增大，内存可能已经无法存放这样海量的数据，这时可以通过写入本地 TFRecords 文件，减小模型训练过程中的内存压力：
with uf.MultiProcess():
X, y = ... # 读取数据
# 缓存数据
model.to_tfrecords(
X=X, y=y, sample_weight=None, X_tokenized=None,
tfrecords_file="train.tfrecords",
)
# 边读边训
model.fit_from_tfrecords(
tfrecords_files=["train.tfrecords", ...], # 支持同步从一个或多个TFRecords文件读取
n_jobs=3, # 启动三个线程
batch_size=32, # 以下参数和`fit`函数中参数相同
learning_rate=5e-05,
target_steps=None,
total_steps=-3,
warmup_ratio=0.1,
print_per_secs=1,
save_per_steps=1000,
**kwargs,
)实际上，也就是把 fit 函数中预处理和模型训练的两个步骤给分开。因此如果需要反复使用同一套数据进行训练，通过以上方式处理能节省更多时间。
预训练参数的 ckpt 文件中，每一个变量都有独立的命名和规格，如 ("layer_1/attention/self/query/kernel", [768, 768])。在上文“模型列表”的详细链接中，我们列示了可以直接读取的公开预训练参数，从这些来源下载的预训练参数无需更多处理。但在其他地方获取的预训练参数，可能会存在与本框架中模型命名/规格不一致的情况。
规格不一致时，变量不可读取。但只有命名不一样时，可以通过下面的方法构建映射，将参数读到模型中：
# 初始化模型，触发读取`ckpt`文件，查看哪些变量初始化失败
model.init()
print(model.uninited_vars)
# 人工进行变量名映射，并重新读取预训练参数
print(uf.list_variables(model.init_checkpoint)) # 在打印的结果中找到对应的参数名
model.assignment_map["layer_1/attention/self/query/kernel"] = model.uninited_vars["bert/encoder/layer_1/attention/self/query/kernel"] # 添加映射关系
model.reinit_from_checkpoint() # 重新读取预训练参数
print(model.uninited_vars) # 在打印的结果中看看初始化失败的变量是否已消失ckpt 是 tensorflow 输出的预训练参数，如果希望读取 PyTorch 输出的预训练参数，则稍微繁琐一些，可以通过将参数读到内存中，使用下面的变量赋值的方法实现。
将内存中的矩阵直接赋值给模型变量：
import numpy as np
array = np.array([[0, 1, 2], [3, 4, 5]]) # 使用numpy.Array格式
print(model.global_variables) # 查看所有变量
variable = model.global_variables[5] # 获取变量
model.assign(variable, array) # 赋值
print(model.sess.run(variable)) # 查看是否赋值成功大多数模型的推理/评分只有以下几个参数，非常简单：
# 推理
model.predict(X=X, X_tokenized=None, batch_size=8)
# 评分
model.score(X=X, y=y, sample_weight=None, X_tokenized=None, batch_size=8)与训练不同的是，推理/评分暂时不支持多进程加速和写入 TFRecords。如果需要推理海量数据，可以通过分批处理达成目的。
模型训练中包含的一些可改动的细节以及训练技巧，在 fit 和 fit_from_tfrecords 函数尾部添加参数即可实现。
自神经网络火热开始，逐渐演化出了一众优秀的最优化算法，包括 Gradient Descent (GD)、Momentum、Adaptive Gradient (AdaGrad)、Root Mean Square prop (RMSprop)、Adaptive Moment estimation (Adam) 等。时至今日，最常见的依然是 2018 年 BERT 所使用的 Adam Weight Decay Regularization (AdamW) 算法。但据实验表明，当训练的 batch_size 达到 512 以上时，AdamW 的收敛效率会极速下降，因而有了后来宣称能够支持大容量 batch 收敛，在 76 分钟内完成 BERT 预训练的 Layer-wise Adaptive Moments optimizer for Batching training (LAMB) 算法。这一套算法的表现基本优于 AdamW 或与之相当，推荐尝试。
model.fit(..., optimizer="gd")
model.fit(..., optimizer="adam")
model.fit(..., optimizer="adamw") # 默认adamw
model.fit(..., optimizer="lamb")迁移学习中常见灾难性遗忘问题 (Catastrophic Forgetting)：模型急不可耐地适应新数据，而丢失了预训练中学到的知识。为了对抗这种过拟合，分层学习率是有效的 trick。越靠近输出层，学习率越大，反之亦然。启用方法是增加 layerwise_lr_decay_ratio 参数并设定一个 0 到 1 之间的浮点数。模型参数会通过 .decay_power 找到自己对应的指数，而后计算学习率：
model.fit(..., layerwise_lr_decay_ratio=0.85)
print(model.decay_power) # 可in-place修改在 2020 年前后，对抗式训练因其较好的效果开始从小众变成人尽皆知。与 CV 领域通过 GAN 的范式进行图像生成的对抗式训练不同，NLP 领域的对抗式训练指的是在梯度正方向上添加扰动 (而不是随机扰动)，以增强 embedding 泛化性和鲁棒性的策略。从最经典的 FGSM 开始，对抗式训练算法也在随时间推进，逐渐优化。在 2020 年，微软推出 SMART 算法成为该领域的 SOTA，在 GLUE 榜单榜上有名。但由于该算法只能在单 label 分类的场景使用，且需要调节的参数较多，因而也存在诸多不便之处。
model.fit(..., adversarial="fgm", epsilon=0.5) # FGM
model.fit(..., adversarial="pgd", epsilon=0.05, n_loop=2) # PGD
model.fit(..., adversarial="freelb", epsilon=0.3, n_loop=3) # FreeLB
model.fit(..., adversarial="freeat", epsilon=0.001, n_loop=3) # FreeAT
model.fit(..., adversarial="smart", epsilon=0.01, n_loop=2, prtb_lambda=0.5, breg_miu=0.2, tilda_beta=0.3) # SMART (仅Classifier可用)全称 Regularized Dropout，同样是行之有效的在对抗过拟合的同时增强泛化能力的策略。策略**极其简单，因而在大多数场景都能应用，在于将相同的数据经过不同的 dropout 后得到各自的概率分布，计算双向 KL 散度并加入到损失函数中。这一想法与对比学习、一致性学习颇为类似，但作用的变量是 dropout。
model.fit(..., rdrop=True, alpha=1.0) # alpha是损失项乘子凝聚了对比学习的核心** —— 在编码上推远负例、拉进正例。正例是相同样本在两次前馈中不同 dropout 下的结果，负例是同一 batch 下的其他样本。私以为，对比学习作为当前 sentence embedding 领域的 SOTA，值得作为多任务学习的 trick，加入到有监督的训练任务中。这里我们实现了 SimCSE 原论文中无监督的训练方案：
tau。
model.fit(..., info_nce_loss=True, tau=1.0, alpha=0.05) # alpha是损失项乘子在基于交叉熵损失的分类场景下，动态调节易/难分样本的损失大小，从而使训练将更多的注意力放在难分样本上。是解决类型不平衡问题的绝佳手段，原论文引用次数已超过 1.5w。
gamma，可取任意大于 0 的值； sample_weight 理念重合，使用参数 sample_weight 即可达到调节
X = ["孔雀东南飞", "五里一徘徊", "不知天上宫阙", "今夕是何年"]
y = [2, 0, 1, 1]
alpha_map = {0: 1.8, 1: 1.0, 2: 1.5}
sample_weight = [alpha_map[_y] for _y in y] # 样本权重，上文中有介绍
model.fit(X, y, sample_weight, focal_loss=True, gamma=1.0)导出供部署上线的 PB 文件到指定目录下：
model.export(
"tf_serving/model", # 导出目录
rename_inputs={}, # 重命名输入
rename_outputs={}, # 重命名输出
ignore_inputs=[], # 剪裁输入
ignore_outputs=[], # 裁剪输出
)而后的模型服务化步骤在 cpp、go、java 等多种语言上都能实现，涉及到的后端代码以及部署上线已与本框架无关，这里就不再展示了。
核心的代码架构如下图所示。新的模型类需要在 apps 目录下添加，而建模则在 model 目录下。欢迎一切有效的 pull request。
-
问：有什么提高训练速度的方法吗？
答：首先是几种能立即实施的基础方法：减小 max_seq_length，多 GPU 并行，多进程数据处理，以及梯度累积。在这些之外，可以进一步尝试对输入的数据进行拆分，在训练过程中逐步提高 max_seq_length、batch_size 和 dropout_rate (通过提高拟合速度，缩短整个训练周期)。当然，还有一些在 UNIF 暂时无法实现的功能，可以前往其他 repo 寻求解决方案，包括但不限于混合精度训练、OP融合、使用 Linformer 等时间复杂度小于
$O(N^2)$ 的模型。 -
问：训练时内存不足，该怎么办？
答：首先需要明确是显存溢出还是内存爆炸。如果是显存溢出，则需要降低 batch_size；如果是由于数据体量过于庞大导致的内存爆炸，可以尝试通过
model.to_tfrecords()分批将数据写入 TFRecords 文件，而后清出内存，通过model.fit_from_tfrecords()读取进行训练。 -
问：模型输入有什么限制吗？
答：对于大多数模型来说，没有限制。一条样本，可以是一个字符串，也可以是多个字符串。以 BERT 为例，完全不必局限于一到两个句子的输入，而是可以通过 list 组合多个 segment，e.g.
X = [["文档1句子1", "文档1句子2", "文档1句子3"], ["文档2句子1", "文档2句子2"]]，模型会自动按顺序拼接并添加分隔符。 -
问：如何查看切词结果？
答：通过
model.tokenizer.tokenize(text)可查看切词结果。另外也可通过model.convert(X)查看经过预处理的实际的模型输入。 -
问：如果我想使用自己的切词工具，该怎么做？
答：提前使用自己的工具做好分词，而后在训练和推理时将原先的传入参数由
X改为X_tokenized。例如，原本传入model.fit(X=["黎明与夕阳", ...], ...)，使用自己的分词工具后，改为传入model.fit(X_tokenized=[["黎", "##明", "与", "夕", "##阳"], ...], ...)。此外，还需保证分词在 vocab 文件中存在。 -
问：无意义的 warning 信息太多，该怎么剔除？
答：这是 tensorflow 一直饱受诟病之处，我也与你一同深受困扰。可以试着在运行时屏蔽 WARNING 类型的信息：
python3 train.py | grep "WARNING" -v。
-
经典永不衰：
模型并非越复杂表现越好，在文本生成以外的应用上，BERT 几乎足够让你一招鲜吃遍天。
-
推荐的预训练参数：
优先使用哈工大训的
macbert-base，亲测各项任务上表现都很不错。 -
TinyBERT 搭配 FastBERT 进行二重蒸馏：
TinyBERTClassifier训练完成后使用.to_bert()提取子模型为 BERT，而后使用FastBERTClassifier读取，继续进行提速。 -
实现一些有趣的事情：
用
GPT2LM来生成古诗/小说，用TransformerMT搭建简单的聊天机器人，或组合ELECTRALM和BERTLM进行文本纠错等等。
ChatGPT 的出现让 NLP 领域变得妙不可言。我已退出原就职公司，亦退出这个行业，此仓库因此搁置。虽然只有寥寥的一百多人关注，但站在我面前也是不小的队伍，真诚感谢大家的关注。有问题依然可以提问在 issue 区，我会尽快回答。