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.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 - 獲取鑒權。
加載運行優化后的模型。
除了增加一行
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)