優化TensorFlow模型

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

1. 導入PAI-Blade和其他依賴庫。

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

2. 編寫一個簡單的函數，用于下載待優化的模型和測試數據。

   雖然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()

3. 調用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!

4. 打印優化報告。

   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"           // 加速比。
     },
     // ......
   }

5. 對比優化前后的性能。

   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.

擴展

后續步驟

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