BERT（Bidirectional Encoder Representation from Transformers）是一个预训练的语言表征模型。作为NLP领域近年来重要的突破，BERT模型在多个自然语言处理的任务中取得了最优结果。然而BERT模型存在巨大的参数规模和计算量，因此实际生产中对该模型具有强烈的优化需求。本文主要介绍如何使用Blade优化通过TensorFlow训练的BERT模型。
使用限制
本文使用的环境需要满足以下版本要求：
- 系统环境：Linux系统中使用Python 3.6及其以上版本、CUDA 10.0。
- 框架：TensorFlow 1.15。
- 推理优化工具：Blade 3.16.0及其以上版本。
步骤一：准备工作
本文以PAI在ModelHub中提供的新闻分类模型为例进行优化，该模型能够将新闻的文本内容分类为教育、三农及娱乐等标签。关于新闻分类模型的详细信息请参见新闻分类模型。
- 执行如下命令安装tokenizers库。
- 下载模型，并解压到指定目录。
wget http://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/tutorials/bert_example/nlu_general_news_classification_base.tar.gz
mkdir nlu_general_news_classification_base
tar zxvf nlu_general_news_classification_base.tar.gz -C nlu_general_news_classification_base
- 使用TensorFlow自带的
saved_model_cli命令查看模型的基本信息。saved_model_cli show --dir nlu_general_news_classification_base --all
命令输出如下结果。
MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:
signature_def['serving_default']:
The given SavedModel SignatureDef contains the following input(s):
inputs['input_ids'] tensor_info:
dtype: DT_INT32
shape: (-1, -1)
name: input_ids:0
inputs['input_mask'] tensor_info:
dtype: DT_INT32
shape: (-1, -1)
name: input_mask:0
inputs['segment_ids'] tensor_info:
dtype: DT_INT32
shape: (-1, -1)
name: segment_ids:0
The given SavedModel SignatureDef contains the following output(s):
outputs['logits'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 28)
name: app/ez_dense/BiasAdd:0
outputs['predictions'] tensor_info:
dtype: DT_INT32
shape: (-1)
name: ArgMax:0
outputs['probabilities'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 28)
name: Softmax:0
Method name is: tensorflow/serving/predict
从上述输出可以看出新闻文本分类模型有三个输入Tensor，分别是
input_ids:0、
input_mask:0及
segment_ids:0。三个输出Tensor，分别是
logits、
predictions及
probabilities，其中
predictions下的
ArgMax:0表示最终分类的类别，即后续关注的推理结果。
- 调用
tokenizers，准备测试数据。from tokenizers import BertWordPieceTokenizer
# 从模型目录的vocab.txt文件初始化tokenizer。
tokenizer = BertWordPieceTokenizer('./nlu_general_news_classification_base/vocab.txt')
# 将四条新闻文本组成一个Batch进行编码。
news = [
'确诊病例超1000例墨西哥宣布进入卫生紧急状态。中新网3月31日电综合报道，墨西哥新冠肺炎病例已超过1000例，墨西哥政府30日宣布进入卫生紧急状态，加强相关措施以遏制新冠肺炎疫情蔓延。',
'国家统计局发布的数据显示，8月份，中国制造业采购经理指数（PMI）为50.1%，继续位于临界点以上，低于上月0.3个百分点。',
'北京时间8月31日讯，在刚刚结束的东京残奥会盲人男足小组赛最后一轮中，中国队依靠朱瑞铭的梅开二度2-0战胜东道主日本，以两胜一负的战绩晋级半决赛。',
'截至8月30日，“祝融号”火星车已在火星表面行驶达100天。100天里，“祝融号”在着陆点以南方向累计行驶1064米，搭载6台科学载荷，共获取约10GB原始科学数据。',
]
tokenized = tokenizer.encode_batch(news)
# 将序列长度填充到128。
def pad(seq, seq_len, padding_val):
return seq + [padding_val] * (seq_len - len(seq))
input_ids = [pad(tok.ids, 128, 0) for tok in tokenized]
segment_ids = [pad(tok.type_ids, 128, 0) for tok in tokenized]
input_mask = [ pad([1] * len(tok.ids), 128, 0) for tok in tokenized ]
# 最终的测试数据是TensorFlow的Feed Dict形式。
test_data = {
"input_ids:0": input_ids,
"segment_ids:0": segment_ids,
"input_mask:0": input_mask,
}
- 加载模型并使用测试数据进行推理。
import tensorflow.compat.v1 as tf
import json
# 加载标签映射文件，获得输出类别整数对应的类别名称。
with open('./nlu_general_news_classification_base/label_mapping.json') as f:
MAPPING = {v: k for k, v in json.load(f).items()}
# 加载并执行模型。
cfg = tf.ConfigProto()
cfg.gpu_options.allow_growth = True
with tf.Session(config=cfg) as sess:
tf.saved_model.loader.load(sess, ['serve'], './nlu_general_news_classification_base')
result = sess.run('ArgMax:0', test_data)
print([MAPPING[r] for r in result])
推理结果如下所示，符合预期。
['国际', '财经', '体育', '科学']
步骤二：调用Blade优化模型
- 调用
blade.optimize对模型进行优化，示例代码如下。关于该接口的详细，请参见Python接口文档。import blade
saved_model_dir = 'nlu_general_news_classification_base'
optimized_model, _, report = blade.optimize(
saved_model_dir, # 模型路径。
'o1', # O1无损优化。
device_type='gpu', # 面向GPU设备优化。
test_data=[test_data] # 测试数据。
)
优化模型时，您需要注意以下事宜：
blade.optimize的第一个返回值为优化后的模型，其数据类型与输入的模型相同。在这个示例中，输入的是SavedModel的路径，返回的是优化后的SavedModel路径。
- 您无需提供
inputs和outputs两个参数，因为Blade可以对输入和输出节点进行自动推断。
- 优化完成后，打印优化报告。
print("Report: {}".format(report))
打印的优化报告类似如下输出。
Report: {
"software_context": [
{
"software": "tensorflow",
"version": "1.15.0"
},
{
"software": "cuda",
"version": "10.0.0"
}
],
"hardware_context": {
"device_type": "gpu",
"microarchitecture": "T4"
},
"user_config": "",
"diagnosis": {
"model": "nlu_general_news_classification_base",
"test_data_source": "user provided",
"shape_variation": "dynamic",
"message": "",
"test_data_info": "input_ids:0 shape: (4, 128) data type: int64\nsegment_ids:0 shape: (4, 128) data type: int64\ninput_mask:0 shape: (4, 128) data type: int64"
},
"optimizations": [
{
"name": "TfStripUnusedNodes",
"status": "effective",
"speedup": "na",
"pre_run": "na",
"post_run": "na"
},
{
"name": "TfStripDebugOps",
"status": "effective",
"speedup": "na",
"pre_run": "na",
"post_run": "na"
},
{
"name": "TfAutoMixedPrecisionGpu",
"status": "effective",
"speedup": "1.46",
"pre_run": "35.04 ms",
"post_run": "24.02 ms"
},
{
"name": "TfAicompilerGpu",
"status": "effective",
"speedup": "2.43",
"pre_run": "23.99 ms",
"post_run": "9.87 ms"
}
],
"overall": {
"baseline": "35.01 ms",
"optimized": "9.90 ms",
"speedup": "3.54"
},
"model_info": {
"input_format": "saved_model"
},
"compatibility_list": [
{
"device_type": "gpu",
"microarchitecture": "T4"
}
],
"model_sdk": {}
}
从优化报告可以看出本示例的优化中
TfAutoMixedPrecisionGpu和
TfAicompilerGpu两个优化项生效，共计带来了3.54倍的加速，将模型推理时间从35 ms提升到了9.9 ms。上述优化结果仅为本示例的测试结果，您的优化效果以实际为准。关于优化报告的字段详情请参见
优化报告。
- 打印
optimized_model的路径。print("Optimized model: {}".format(optimized_model))
系统输出如下类似结果。
Optimized model: /root/nlu_general_news_classification_base_blade_opt_20210901141823/nlu_general_news_classification_base
从上述输出结果可以看出优化后的模型已经存放在新的路径下了。
步骤三：验证性能与正确性
优化完成后，通过Python脚本对优化报告的信息进行验证。
- 定义
benchmark方法，对模型进行10次预热，然后运行1000次，最终取平均的推理时间作为推理速度。import time
def benchmark(model, test_data):
tf.reset_default_graph()
with tf.Session() as sess:
sess.graph.as_default()
tf.saved_model.loader.load(sess, ['serve'], model)
# Warmup!
for i in range(0, 10):
result = sess.run('ArgMax:0', test_data)
# Benchmark!
num_runs = 1000
start = time.time()
for i in range(0, num_runs):
result = sess.run('ArgMax:0', test_data)
elapsed = time.time() - start
rt_ms = elapsed / num_runs * 1000.0
# Show the result!
print("Latency of model: {:.2f} ms.".format(rt_ms))
print("Predict result: {}".format([MAPPING[r] for r in result]))
- 调用
benchmark方法，对原始模型进行验证。benchmark('nlu_general_news_classification_base', test_data)
系统返回如下类似结果。
Latency of model: 36.20 ms.
Predict result: ['国际', '财经', '体育', '科学']
从结果可以看出推理时间36.20 ms与优化报告中"overall"下的
"baseline": "35.01 ms"基本一致。预测结果
['国际', '财经', '体育', '科学']与预期的结果一致。此处的推理时间仅为本案例的测试结果，您模型的推理时间以实际结果为准。
- 调用
benchmark方法，对优化后的型进行验证。import os
os.environ['TAO_COMPILATION_MODE_ASYNC'] = '0'
benchmark(optimized_model, test_data)
由于优化报告显示AICompiler对模型产生了优化效果，而AICompiler是异步编译的，即在编译过程中仍然会使用原有的模型进行推理。因此，为了测试数据的准确性，在调用
benchmark前，需要设置环境变量
TAO_COMPILATION_MODE_ASYNC=0强制地将编译设置为同步模式。
系统返回如下类似结果。
Latency of model: 9.87 ms.
Predict result: ['国际', '财经', '体育', '科学']
从结果可以看出推理时间9.87 ms与优化报告中"overall"下的
"optimized": "9.90 ms"基本一致。预测结果
['国际', '财经', '体育', '科学']与预期的结果一致。此处的推理时间仅为本案例的测试结果，您模型的推理时间以实际结果为准。
步骤四：加载运行优化后的模型
完成验证后，您需要对模型进行部署，Blade提供了Python和C++两种运行时SDK供您集成。关于C++的SDK使用方法请参见使用SDK部署TensorFlow模型推理，下文主要介绍如何使用Python SDK部署模型。
- 可选：在试用阶段，您可以设置如下的环境变量，防止因为鉴权失败而程序退出。
export BLADE_AUTH_USE_COUNTING=1
- 获取鉴权。
export BLADE_REGION=<region>
export BLADE_TOKEN=<token>
您需要根据实际情况替换以下参数：
- <region>：Blade支持的地域，需要加入Blade用户群获取该信息，用户群的二维码详情请参见获取Token。
- <token>：鉴权Token，需要加入Blade用户群获取该信息，用户群的二维码详情请参见获取Token。
- 加载运行优化后的模型。
除了增加一行
import blade.runtime.tensorflow，您无需为Blade的接入编写额外代码，即原有的推理代码无需任何改动。
import tensorflow.compat.v1 as tf
import blade.runtime.tensorflow
# <your_optimized_model_path>替换为优化后的模型路径。
savedmodel_dir = <your_optimized_model_path>
# <your_infer_data>替换为用于推理的数据。
infer_data = <your_infer_data>
with tf.Session() as sess:
sess.graph.as_default()
tf.saved_model.loader.load(sess, ['serve'], savedmodel_dir)
result = sess.run('ArgMax:0', infer_data)