最佳化TensorFlow模型

PAI-Blade提供了豐富的模型最佳化方法，您只需要在本地環境中安裝Wheel包，即可通過調用Python API的方式進行模型最佳化。本文介紹如何使用PAI-Blade最佳化TensorFlow模型，所有實驗結果均在NVidia T4卡上測得。

前提條件

已安裝TensorFlow及PAI-Blade的Wheel包，詳情請參見安裝Blade。
已有訓練完成的TensorFlow模型，本文使用一個公開的ResNet50模型。

本文以一個公開的ResNet50模型為例，示範如何最佳化TensorFlow模型。您也可以對自己的TensorFlow模型進行最佳化。

匯入PAI-Blade和其他依賴庫。

import os
import numpy as np
import tensorflow.compat.v1 as tf
import blade

編寫一個簡單的函數，用於下載待最佳化的模型和測試資料。

雖然PAI-Blade支援沒有測試資料的最佳化，即零輸入最佳化，但是基於真實輸入資料的最佳化結果會更加準確有效。因此，建議您提供測試資料。下載待最佳化的模型和測試資料的函數樣本如下所示。

def _wget_demo_tgz():
    # 下載一個公開的resnet50模型。
    url = 'http://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/demo/mask_rcnn_resnet50_atrous_coco_2018_01_28.tar.gz'
    local_tgz = os.path.basename(url)
    local_dir = local_tgz.split('.')[0]
    if not os.path.exists(local_dir):
        blade.util.wget_url(url, local_tgz)
        blade.util.unpack(local_tgz)
    model_path = os.path.abspath(os.path.join(local_dir, "frozen_inference_graph.pb"))
    graph_def = tf.GraphDef()
    with open(model_path, 'rb') as f:
        graph_def.ParseFromString(f.read())
    # 以隨機數作為測試資料。
    test_data = np.random.rand(1, 800,1000, 3)
    return graph_def, {'image_tensor:0': test_data}

graph_def, test_data = _wget_demo_tgz()

調用blade.optimize函數進行模型最佳化，詳細的參數解釋請參見Python介面文檔。模型最佳化的程式碼範例如下。

input_nodes=['image_tensor']
output_nodes = ['detection_boxes', 'detection_scores', 'detection_classes', 'num_detections', 'detection_masks']

optimized_model, opt_spec, report = blade.optimize(
    graph_def,                 # 待最佳化的模型，此處是tf.GraphDef, 也可以配置為SavedModel的路徑。
    'o1',                      # 最佳化層級，o1或o2。
    device_type='gpu',         # 目標裝置，gpu/cpu/edge。
    inputs=input_nodes,        # 輸入節點，也可以不提供，PAI-Blade會自行推斷。
    outputs=output_nodes,      # 輸出節點。
    test_data=[test_data]      # 測試資料。
)

blade.optimize函數返回的三個對象，分別如下所示：

optimized_model：最佳化完成的模型，本文為tf.GraphDef。
opt_spec：包含複現最佳化結果需要的配置資訊、環境變數及資源檔等，通過with語句可以使其生效。
report：最佳化報告，可以直接列印。關於報告中的參數解釋，詳情請參見最佳化報告。

最佳化過程中，您可以看到如下類似的最佳化進度。

[Progress] 5%, phase: user_test_data_validation.
[Progress] 10%, phase: test_data_deduction.
[Progress] 15%, phase: CombinedSwitch_1.
[Progress] 24%, phase: TfStripUnusedNodes_22.
[Progress] 33%, phase: TfStripDebugOps_23.
[Progress] 42%, phase: TfFoldConstants_24.
[Progress] 51%, phase: CombinedSequence_7.
[Progress] 59%, phase: TfCudnnrnnBilstm_25.
[Progress] 68%, phase: TfFoldBatchNorms_26.
[Progress] 77%, phase: TfNonMaxSuppressionOpt_27.
[Progress] 86%, phase: CombinedSwitch_20.
[Progress] 95%, phase: model_collecting.
[Progress] 100%, Finished!

列印最佳化報告。

print("Report: {}".format(report))

在最佳化報告中可以看到主要的效果源於哪些最佳化項，如下所示。

Report: {
  // ......
  "optimizations": [
    // ......
    {
      "name": "TfNonMaxSuppressionOpt",
      "status": "effective",
      "speedup": "1.58",        // 加速比。
      "pre_run": "522.74 ms",   // 最佳化前延遲。
      "post_run": "331.45 ms"   // 最佳化後延遲。
    },
    {
      "name": "TfAutoMixedPrecisionGpu",
      "status": "effective",
      "speedup": "2.43",
      "pre_run": "333.30 ms",
      "post_run": "136.97 ms"
    }
    // ......
  ],
  // 端到端最佳化結果。
  "overall": {
    "baseline": "505.91 ms",    // 原始模型延遲。
    "optimized": "136.83 ms",   // 最佳化後模型延遲。
    "speedup": "3.70"           // 加速比。
  },
  // ......
}

對比最佳化前後的效能。

import time

def benchmark(model):
    tf.reset_default_graph()
    with tf.Session() as sess:
        sess.graph.as_default()
        tf.import_graph_def(model, name="")
        # Warmup!
        for i in range(0, 1000):
            sess.run(['image_tensor:0'], test_data)
        # Benchmark!
        num_runs = 1000
        start = time.time()
        for i in range(0, num_runs):
            sess.run(['image_tensor: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))

# 對原模型測速。
benchmark(graph_def)

# 對最佳化後的模型測速。
with opt_spec:
    benchmark(optimized_model)

效能實測結果如下，該結果與最佳化報告中的值基本一致。

Latency of model: 530.26 ms.
Latency of model: 148.40 ms.

擴充

blade.optimize函數的model參數支援多種形式的模型輸入。對於TensorFlow模型，支援以下三種方式傳入模型：

直接傳入tf.GraphDef對象
從檔案載入PB或PBTXT格式的Frozen PB
從指定路徑匯入SavedModel

在本文樣本中，為blade.optimize函數傳入了記憶體中的tf.GraphDef對象。另外兩種方式可以參考如下代碼：

傳入Fronzen PB檔案

optimized_model, opt_spec, report = blade.optimize(
    './path/to/frozen_pb.pb',  # 也可以為.pbtxt格式。
    'o1',
    device_type='gpu',
)

傳入SavedModel路徑

optimized_model, opt_spec, report = blade.optimize(
    './path/to/saved_model_directory/',
    'o1',
    device_type='gpu',
)

後續步驟

經過PAI-Blade最佳化的模型，您可以通過Python直接執行或部署為EAS服務。此外，PAI-Blade也提供了C++ SDK，以便您將最佳化後的模型整合到自己的應用中，詳情請參見使用SDK部署TensorFlow模型推理。