中文 | English
有数据,想要快速实现你的想法?轻便、易使用的自然语言处理联合框架,帮你快速搭建各类常用深度学习模型 (Transformer, GPT-2, BERT, ALBERT, UniLM, XLNet, ELECTRA),同时对于 BERT 系列,支持高效用的蒸馏 (TinyBERT, FastBERT)。支持各类上下游任务 (语言模型、文本分类、文本生成、命名实体识别、机器阅读理解、机器翻译、序列标注等)。
- API 简单:三行代码完成训练及推理,并一键设置多 GPU 并行
- 品类丰富:支持 40+ 种模型类
- 唯一依赖:Tensorflow 1.x/2.x
- 高分保证:提供分层学习率、对抗式训练等多项训练技巧
- 高扩展性:代码高度抽象、易扩展,能够高效支持更多 SOTA 算法
- 可供部署:导出 SavedModel,供部署使用
git clone https://github.com/geyingli/unif
cd unif
python3 setup.py install --user
import uf
# 载入模型(使用 demo 配置文件进行示范)
model = uf.BERTClassifier(config_file='demo/bert_config.json', vocab_file='demo/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 建模的全能语言模型 | |
命名实体识别 | BERTNER |
- |
BERTCRFNER |
结合 CRF | |
BERTCRFCascadeNER |
识别与分类同时进行的级联架构 | |
机器翻译 | TransformerMT |
共享词表,标准 Seq2Seq 架构 |
机器阅读理解 | BERTMRC |
- |
RoBERTaMRC |
- | |
ALBERTMRC |
- | |
ELECTRAMRC |
- | |
SANetMRC |
引入 Sentence Attention | |
BERTVerifierMRC |
抽取 answer span 的同时判断可答性 | |
RetroReaderMRC |
抽取 answer span 的同时判断可答性 | |
单 Label 分类 | TextCNNClassifier |
小而快 |
BERTClassifier |
- | |
XLNetClassifier |
- | |
RoBERTaClassifier |
- | |
ALBERTClassifier |
- | |
ELECTRAClassifier |
- | |
WideAndDeepClassifier |
通过 Wide & Deep 架构融合句子级别特征 | |
SemBERTClassifier |
通过 SemBERT 架构融合字级别的特征 | |
PerformerClassifier |
引入 FAVOR+ 加速推理 | |
UDAClassifier |
结合一致性学习的半监督学习算法 | |
多 Label 分类 | BERTBinaryClassifier |
- |
XLNetBinaryClassifier |
- | |
RoBERTaBinaryClassifier |
- | |
ALBERTBinaryClassifier |
- | |
ELECTRABinaryClassifier |
- | |
序列标注 | BERTSeqClassifier |
- |
XLNetSeqClassifier |
- | |
RoBERTaSeqClassifier |
- | |
ALBERTSeqClassifier |
- | |
ELECTRASeqClassifier |
- | |
模型蒸馏 | TinyBERTClassifier |
大幅压缩模型参数,提速十倍以上 |
TinyBERTBinaryClassifier |
- | |
FastBERTClassifier |
动态推理,易分样本提前离开模型 |
文档目前还不完善,善用 help(XXX)
能帮你获得更多 API 的使用细节。
一步创建新模型:
model = uf.BERTClassifier(
config_file, vocab_file,
max_seq_length=128,
label_size=2,
init_checkpoint=None, # 预训练参数路径
output_dir='./output',
gpu_ids='0,1,3,5',
drop_pooler=False, # 建模时跳过 pooler 层
do_lower_case=True,
truncate_method='LIFO') # longer-FO/LIFO/FIFO
下载知名公开预训练参数:
# 查看可下载列表
uf.list_resources()
┌──────────────────────────┬──────────┬──────────────┬───────────────────────────────────────────┬───────────────────────────────────────────────────────────────────────────────────┐
┊ Key ┊ Backbone ┊ Organization ┊ Site ┊ URL ┊
├──────────────────────────┼──────────┼──────────────┼───────────────────────────────────────────┼───────────────────────────────────────────────────────────────────────────────────┤
┊ bert-base-zh ┊ BERT ┊ Google ┊ https://github.com/google-research/bert ┊ https://storage.googleapis.com/bert_models/2018_11_03/chinese_L-12_H-768_A-12.zip ┊
┊ albert-tiny-zh ┊ ALBERT ┊ Google ┊ https://github.com/google-research/albert ┊ https://storage.googleapis.com/albert_zh/albert_tiny_zh_google.zip ┊
...
┊ electra-180g-large-zh ┊ ELECTRA ┊ HFL ┊ https://github.com/ymcui/Chinese-ELECTRA ┊ https://drive.google.com/uc?export=download&id=1P9yAuW0-HR7WvZ2r2weTnx3slo6f5u9q ┊
└──────────────────────────┴──────────┴──────────────┴───────────────────────────────────────────┴───────────────────────────────────────────────────────────────────────────────────┘
# 下载预训练模型包
uf.download('bert-wwm-ext-base-zh')
任务后期需要大量的训练,可以通过配置文件,方便地整理和读取模型:
# 写入配置文件
assert model.output_dir is not None # 为空的话模型就白训了
model.cache('key', cache_file='.cache')
# 从配置文件读取
model = uf.load('key', cache_file='.cache')
程序还没执行结束,内存就不够用了?试试删除模型 del model
或重置 model.reset()
。
# 开启多进程 (加速数据处理)
with uf.MultiProcess(): # 建议在这一步之后再读取大批量数据,否则容易内存爆炸
# 训练
model.fit(
X=X, y=y, sample_weight=None,
X_tokenized=None, # 特定场景下使用,e.g. 使用你自己的分词工具/语言模型推理时在输入加 [MASK]
batch_size=32,
learning_rate=5e-05,
target_steps=None, # 放空代表直接训练到 `total_steps`,不中途停止;否则为本次训练暂停点
total_steps=-3, # -3 代表自动计算数据量并循环三轮
warmup_ratio=0.1,
print_per_secs=1, # 多少秒打印一次信息
save_per_steps=1000,
**kwargs) # 其他参数,下文介绍
# 推理
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)
# 常规训练流程示范
assert model.output_dir is not None # 非空才能保存模型参数
for loop_id in range(10): # 假设训练途中一共验证 10 次
model.fit(X, y, target_steps=((loop_id + 1) * -0.6), total_steps=-6) # 假设一共训练 6 轮
model.cache('dev-%d' % loop_id) # 保存一次模型
print(model.score(X_dev, y_dev)) # 查看模型表现
复用训练数据?可以尝试先存为 TFRecords,训练时读取:
with uf.MultiProcess():
# 缓存数据
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-0', './.tfrecords-1'], # 同时从两个 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)
训练所用的条件参数 kwargs:
# 优化器
model.fit(X, y, ..., optimizer='gd')
model.fit(X, y, ..., optimizer='adam')
model.fit(X, y, ..., optimizer='adamw') # 默认
model.fit(X, y, ..., optimizer='lamb')
# 分层学习率 (少量模型不适用)
model.fit(X, y, ..., layerwise_lr_decay_ratio=0.85) # 默认为 None
print(model._key_to_depths) # 衰减比率
# 对抗式训练
model.fit(X, y, ..., adversarial='fgm', epsilon=0.5) # FGM
model.fit(X, y, ..., adversarial='pgd', epsilon=0.05, n_loop=2) # PGD
model.fit(X, y, ..., adversarial='freelb', epsilon=0.3, n_loop=3) # FreeLB
model.fit(X, y, ..., adversarial='freeat', epsilon=0.001, n_loop=3) # FreeAT
model.fit(X, y, ..., adversarial='smart', epsilon=0.01, n_loop=2, prtb_lambda=0.5, breg_miu=0.2, tilda_beta=0.3) # SMART (仅 Classifier 可用)
# 置信度过滤 (仅 Classifier 可用)
model.fit(X, y, ..., conf_thresh=0.99) # 默认为 None
存在变量命名不同而无法加载,可通过以下步骤解决:
# 查看从 `init_checkpoint` 初始化失败的变量
assert model.init_checkpoint is not None
model.init()
print(model.uninited_vars)
# 在 `checkpoint` 中寻找对应的参数名
print(uf.list_variables(model.init_checkpoint))
# 人工添加映射关系到 `assignment_map`
model.assignment_map['var_1_in_ckpt'] = model.uninited_vars['var_1_in_model']
model.assignment_map['var_2_in_ckpt'] = model.uninited_vars['var_2_in_model']
# 重新读取预训练参数
model.reinit_from_checkpoint()
# 看看变量是否从初始化失败的名单中消失
print(model.uninited_vars)
# 保存参数及配置(避免下次载入预训练参数时,重复上述步骤)
assert model.output_dir is not None
model.cache('key')
直接给参数赋值如何?当然是可以的:
import numpy as np
# 获取参数
variable = model.trainable_variables[0]
# 赋值
model.assign(variable, value)
# 查看参数
print(model.sess.run(variable))
# 保存赋值后的参数及配置
assert model.output_dir is not None
model.cache('key')
# 导出 PB 文件到 `output_dir` 下
assert model.output_dir is not None
model.export(
export_dir, # 导出目录
rename_inputs={}, # 重命名输入
rename_outputs={}, # 重命名输出
ignore_inputs=[], # 剪裁多余输入
ignore_outputs=[]) # 裁剪多余输出
-
问:如何实现多个 segment 的输入?
答:使用 list 组合多个 segment 的输入,如
X = [['文档1句子1', '文档1句子2', '文档1句子3'], ['文档2句子1', '文档2句子2']]
,模型会自动按顺序拼接并添加分隔符。 -
问:如何查看切词结果?
答:通过
model.tokenizer.tokenize(text)
可查看切词结果。另外也可通过model.convert(X)
查看切词与 ID 转换后的矩阵。 -
问:如何使用自己的切词工具?
答:在训练和推理时预先将传入参数
X
改为X_tokenized
,模型将直接跳过原有的的分词步骤。需要注意的是,分词结果同样需要基于list
承载,例如原先由x
承载的['黎明与夕阳']
,由X_tokenized
承载后需呈现['黎', '##明', '与', '夕', '##阳']
的形式。 -
问:如何实现 TinyBERT 和 FastBERT 复蒸馏?
答:
TinyBERTClassifier
训练完成后使用.to_bert()
将变量重命名保存,而后使用FastBERTClassifier
常规读取生成的 checkpoint 和配置文件即可。
我们欢迎一切有效的 pull request,加入我们,成为 contributors 一员。 核心的代码架构如下图所示,新的模型开发仅需要在 application 下添加新的类,这些类可以由现有算法库 modeling 中的算子组合而来,也可以自行编写。
beta v2.5.0 -> Now:
- 新增:添加
uf.TinyBERTBinaryClassifier
,使用TinyBERT架构蒸馏BERT多分类模型 - 新增:添加
uf.MultiProcess
,可以通过多进程加速模型前的数据解析流程 - 新增:添加
uf.list_resources()
和uf.download()
用法,方便地下载公开预训练参数 - 新增:使用
model.cache()
缓存模型时,可以传入note
参数作为备注,传入max_to_keep
参数限制 checkpoint 数量 - 新增:添加完整的示范代码,在 unif/examples 目录下
- 变更:取消训练结束自动保存模型参数,用户需要通过
model.save()
或model.cache()
自行保存 - 变更:默认打印所有信息,无需用户再行输入
uf.set_verbosity()
- 变更:在 vocab 中添加新的 token 时进行提示
- 变更:将优化器动量从
model.uninited_vars
中删除 - 变更:加载预训练参数时,跳过形状不一致的参数,避免报错
- 变更:添加语言模型采样
[MASK]
进度提示 - 变更:
FastBERTClassifier
设置为默认跳过第一个子分类器 - 变更:
UniLM
在s2s
模式下的结尾 token 从[SEP]
改为[EOS]
- 变更:
VAELM
取消返回[CLS]
- 修复:修复了
RoBERTaLM
相关的若干 bug - 修复:修复了多GPU训练及对抗式学习的若干 bug
框架目前主要为我个人及团队所用,靠兴趣推动至今。如果能受到更多人,包括您的认可,我们会愿意投入更多精力进行丰富与完善,早日推出第一个正式版本。如果您喜欢,请点个 star 作为支持。如果有希望实现的 SOTA 算法,留下 issue,我们会酌情考虑,并安排时间为你编写。通常三日以内可以实现。任何需求和建议都不尽欢迎。最后,感谢你读到这里。