如何使用Blade优化通过TensorFlow训练的BERT模型_人工智能平台 PAI(PAI)-阿里云帮助中心

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及其以上版本。

操作流程

使用Blade优化BERT模型的流程如下：

步骤一：准备工作
下载模型，并使用tokenizers库准备测试数据。
步骤二：调用Blade优化模型
调用blade.optimize接口优化模型，并保存优化后的模型。
步骤三：验证性能与正确性
对优化前后的推理速度及推理结果进行测试，从而验证优化报告中信息的正确性。
步骤四：加载运行优化后的模型
集成Blade SDK，加载优化后的模型进行推理。

步骤一：准备工作

执行如下命令安装tokenizers库。
```
pip3 install 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)