如何使用Blade最佳化通過TensorFlow訓練的BERT模型 - Platform For AI

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)