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模型的流程如下:
  1. 步骤一:准备工作

    下载模型,并使用tokenizers库准备测试数据。

  2. 步骤二:调用Blade优化模型

    调用blade.optimize接口优化模型,并保存优化后的模型。

  3. 步骤三:验证性能与正确性

    对优化前后的推理速度及推理结果进行测试,从而验证优化报告中信息的正确性。

  4. 步骤四:加载运行优化后的模型

    集成Blade SDK,加载优化后的模型进行推理。

步骤一:准备工作

本文以PAI在ModelHub中提供的新闻分类模型为例进行优化,该模型能够将新闻的文本内容分类为教育、三农及娱乐等标签。关于新闻分类模型的详细信息请参见新闻分类模型

  1. 执行如下命令安装tokenizers库。
    !pip3 install tokenizers
  2. 下载模型,并解压到指定目录。
    !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
  3. 使用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:0input_mask:0segment_ids:0。三个输出Tensor,分别是logitspredictionsprobabilities,其中predictions下的ArgMax:0表示最终分类的类别,即后续关注的推理结果。
  4. 调用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,
    }
  5. 加载模型并使用测试数据进行推理。
    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优化模型

  1. 调用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路径。
    • 您无需提供inputsoutputs两个参数,因为Blade可以对输入和输出节点进行自动推断。
  2. 优化完成后,打印优化报告。
    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": {}
    }
    从优化报告可以看出本示例的优化中TfAutoMixedPrecisionGpuTfAicompilerGpu两个优化项生效,共计带来了3.54倍的加速,将模型推理时间从35 ms提升到了9.9 ms。上述优化结果仅为本示例的测试结果,您的优化效果以实际为准。关于优化报告的字段详情请参见优化报告
  3. 打印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脚本对优化报告的信息进行验证。

  1. 定义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]))
  2. 调用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"基本一致。预测结果['国际', '财经', '体育', '科学']与预期的结果一致。此处的推理时间仅为本案例的测试结果,您模型的推理时间以实际结果为准。
  3. 调用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部署模型。

  1. 可选:在试用阶段,您可以设置如下的环境变量,防止因为鉴权失败而程序退出。
    export BLADE_AUTH_USE_COUNTING=1
  2. 获取鉴权。
    export BLADE_REGION=<region>
    export BLADE_TOKEN=<token>
    您需要根据实际情况替换以下参数:
    • <region>:Blade敏捷版支持的地域,需要加入Blade用户群获取该信息,用户群的二维码详情请参见获取Token
    • <token>:鉴权Token,需要加入Blade用户群获取该信息,用户群的二维码详情请参见获取Token
  3. 加载运行优化后的模型。
    除了增加一行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)