AI編譯器最佳化 - Platform For AI

AICompiler是整合在PAI-Blade中的AI編譯最佳化組件，包含Static Shape和Dynamic Shape編譯架構。通常您無需提供額外配置，AICompiler即可在通用透明的情況下協助您提高推理效能。本文介紹如何使用AICompiler對TensorFlow和PyTorch模型進行編譯最佳化。

背景資訊

隨著AI模型結構的快速演化，底層計算硬體的層出不窮，使用者使用習慣的推陳出新，使得單純基於手工最佳化解決AI模型的效能和效率問題越來越困難。為瞭解決這些問題，AI編譯最佳化技術已經成為一個被廣泛關注的技術方向。

傳統編譯器是以高層語言作為輸入，避免使用者直接寫機器碼。深度學習編譯器的作用與其相仿，輸入是比較靈活的、具備較高抽象度的計算圖，輸出包括CPU或GPU等硬體平台上的底層機器碼及執行引擎。AI編譯器的目標是針對AI計算任務，以通用編譯器的方式完成效能最佳化。讓使用者可以專註於上層模型開發，降低使用者手工最佳化效能的人力開發成本，進一步壓榨硬體效能空間。

在過去兩年多時間裡，PAI團隊在AI編譯最佳化技術方向投入了比較專註的資源精力，AICompiler已經作為最佳化組件之一整合到PAI-Blade之中，從而協助使用者以通用透明的方式完成推理模型的最佳化和部署。

目前AICompiler主要包含Static Shape和Dynamic Shape兩套編譯架構。Dynamic Shape可以支援所有任務類型，包括計算圖Shape變化範圍較大的推理任務。Static Shape主要適配計算圖為靜態Shape的任務，或Shape變化範圍較小的任務，以獲得理論上更加極致的效能。Static Shape和Dynamic Shape的關係如下表所示。

編譯架構	TensorFlow	PyTorch	適配範圍
Static Shape	支援	暫不支援	適配靜態Shape或Shape變化範圍較小的任務，理論上可以獲得極致的效能收益。
Dynamic Shape	支援	支援	適配所有的任務類型。

預設情況下，PAI-Blade會自動分析使用者的模型是否適合Dynamic Shape的編譯最佳化，使用者無需提供額外的輸入資訊。對於靜態Shape類任務或Shape變化範圍較小的任務，使用者可以提供額外的輸入選項，此時PAI-Blade會嘗試通過Static Shape模式編譯以獲得更優的效能。下文詳細介紹幾個應用例子供您參考。

通過Dynamic Shape模式編譯TensorFlow模型

以一個開源ASR模型為例，示範如何通過Dynamic Shape模式編譯TensorFlow模型：

下載模型和測試資料。

# 下載樣本模型、測試資料。
wget https://pai-blade.cn-hangzhou.oss.aliyun-inc.com/test_public_model/bbs/tf_aicompiler_demo/frozen.pb
wget https://pai-blade.cn-hangzhou.oss.aliyun-inc.com/test_public_model/bbs/tf_aicompiler_demo/test_bc4.npy

載入模型、指定資料，並調用blade.optimize，無需任何額外配置。

import numpy as np
import tensorflow as tf
import blade

# 載入模型和測試資料。
graph_def = tf.GraphDef()
with open('./frozen.pb', 'rb') as f:
    graph_def.ParseFromString(f.read())
test_data = np.load('test_bc4.npy', allow_pickle=True, encoding='bytes',).item()

# 開始最佳化。
optimized_model, opt_spec, report = blade.optimize(
    graph_def,  # The original model, here is a TF GraphDef.
    'o1',  # Optimization level o1 or o2.
    device_type='gpu',  # Target device to run the optimized model.
    config=blade.Config(),
    inputs=['encoder_memory_placeholder', 'encoder_memory_length_placeholder'],
    outputs=['score', 'seq_id'],
    test_data=[test_data]
    #verbose=True
)

# 儲存最佳化結果。
tf.train.write_graph(optimized_model, "./", "optimized.pb", as_text=False)
print("Report: {}".format(report))

最佳化完成後，查看最佳化報告（blade.optimize返回的report）中AICompiler生效之後的效能收益。本例中以T4卡為例，AICompiler可以幫您在通用透明的情況下獲得2.23倍的效能收益，report中的欄位含義請參見最佳化報告。
```
    {
      "name": "TfAicompilerGpu",
      "status": "effective",
      "speedup": "2.23",
      "pre_run": "120.54 ms",
      "post_run": "53.99 ms"
    }
```

通過Static Shape模式編譯TensorFlow模型

如果模型部署後的實際Shape變化範圍較小或完全是Static Shape，您可以通過配置config選項，指定PAI-Blade以Static Shape模式進行編譯。相比Dynamic Shape模式編譯，您只需修改blade.optimize函數的config入參即可實現Static Shape編譯，範例程式碼如下。

optimized_model, opt_spec, report = blade.optimize(
    graph_def,  # The original model, here is a TF GraphDef.
    'o1',  # Optimization level o1 or o2.
    device_type='gpu',  # Target device to run the optimized model.
    # Provide an additional config here in order to try Static Shape Compilation:
    config=blade.Config(enable_static_shape_compilation_opt = True), 
    inputs=['encoder_memory_placeholder', 'encoder_memory_length_placeholder'],
    outputs=['score', 'seq_id'],
    test_data=[test_data]
    #verbose=True
)

關於config參數的進階配置，詳情請參見blade.Config。

最佳化完成後，可以查看最佳化報告（blade.optimize返回的report）中的效能收益。本例中的效能收益進一步提升至2.35倍，report中的欄位含義請參見最佳化報告。

    {
      "name": "TfAicompilerGpu",
      "status": "effective",
      "speedup": "2.35",
      "pre_run": "114.91 ms",
      "post_run": "48.86 ms"
    }

通過Dynamic Shape模式編譯PyTorch模型

以一個開源ASR模型為例，示範如何通過Dynamic Shape模式編譯PyTorch模型：

下載模型。

# PyTorch 1.6.0
# Python3.6
wget https://pai-blade.cn-hangzhou.oss.aliyun-inc.com/test_public_model/bbs/pt_aicompiler_demo/orig_decoder_v2.pt

載入模型、指定資料，並調用blade.optimize，無需任何額外配置。

import os
import time
import torch
# To use blade, just import it.
import blade

# 載入模型。
pt_file = 'orig_decoder_v2.pt'
batch = 8
model = torch.jit.load(pt_file)

# 準備測試資料。
def get_test_data(batch_size=1):
    decoder_input_t = torch.LongTensor([1] * batch_size).cuda()
    decoder_hidden_t = torch.rand(batch_size, 1, 256).cuda()
    decoder_hidden_t = decoder_hidden_t * 1.0
    decoder_hidden_t = torch.tanh(decoder_hidden_t)
    output_highfeature_t = torch.rand(batch_size, 448, 4, 50).cuda()
    attention_sum_t = torch.rand(batch_size, 1, 4, 50).cuda()
    decoder_attention_t = torch.rand(batch_size, 1, 4, 50).cuda()
    et_mask = torch.rand(batch_size, 4, 50).cuda()

    return (decoder_input_t, decoder_hidden_t, output_highfeature_t, attention_sum_t, decoder_attention_t, et_mask)

dummy = get_test_data(batch)

# 開始最佳化。
optimized_model, opt_spec, report = blade.optimize(
    model,  # The original model, here is a torch scrip model.
    'o1',  # Optimization level o1 or o2.
    device_type='gpu',  # Target device to run the optimized model.
    test_data=[dummy],  # For PyTorch, input data is list of tupoles.
    config=blade.Config()
)

print("spec: {}".format(opt_spec))
print("report: {}".format(report))

# 匯出最佳化結果。
torch.jit.save(optimized_model, 'optimized_decoder.pt')

最佳化完成後，查看最佳化報告（blade.optimize返回的report）中AICompiler生效之後的效能收益。本例中以T4卡為例，AICompiler可以幫您在通用透明的情況下獲得2.45倍的效能收益，report中的欄位含義請參見最佳化報告。
```
  "optimizations": [
    {
      "name": "PyTorchMlir",
      "status": "effective",
      "speedup": "2.45",
      "pre_run": "1.99 ms",
      "post_run": "0.81 ms"
    }
  ],
```