最佳化PyTorch模型

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

前提條件

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

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

import os
import time
import torch
import torchvision.models as models
import blade

從torchvision載入ResNet50模型。由於PAI-Blade僅支援ScriptModule，因此需要轉換模型格式。

model = models.resnet50().float().cuda()  # 準備模型。
model = torch.jit.script(model).eval()    # 轉換成ScriptModule。
dummy = torch.rand(1, 3, 224, 224).cuda() # 構造測試資料。

調用blade.optimize函數進行模型最佳化，詳細的參數解釋請參見Python介面文檔。模型最佳化的程式碼範例如下。如果您在最佳化中遇到問題，可以加入Blade客戶群，諮詢相關人員，詳情請參見擷取Token。
```
optimized_model, opt_spec, report = blade.optimize(
    model,                 # 待最佳化的模型。
    'o1',                  # 最佳化層級，o1或o2。
    device_type='gpu',     # 目標裝置，gpu/cpu。
    test_data=[(dummy,)],  # PyTorch的輸入資料是List of tuple of tensor。
)
```
blade.optimize函數返回的三個對象，分別如下所示：
- optimized_model：最佳化完成的模型，此處為torch.jit.ScriptModule。
- opt_spec：包含複現最佳化結果需要的配置資訊、環境變數及資源檔等，通過with語句可以使其生效。
- report：最佳化報告，可以直接列印。關於報告中的參數解釋，詳情請參見最佳化報告。
最佳化過程中，您可以看到如下類似的最佳化進度。
```
[Progress] 5%, phase: user_test_data_validation.
[Progress] 10%, phase: test_data_deduction.
[Progress] 15%, phase: CombinedSwitch_4.
[Progress] 95%, phase: model_collecting.
```

列印最佳化報告。

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

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

Report: {
  // ......
  "optimizations": [
    {
      "name": "PtTrtPassFp32",
      "status": "effective",
      "speedup": "1.50",     // 加速比。
      "pre_run": "5.29 ms",  // 加速前延遲。
      "post_run": "3.54 ms"  // 加速後延遲。
    }
  ],
  // 端到端最佳化結果。
  "overall": {
    "baseline": "5.30 ms",    // 原始模型延遲。
    "optimized": "3.59 ms",   // 最佳化後模型延遲。
    "speedup": "1.48"         // 加速比。
  },
  // ......
}

對比最佳化前後的效能。

@torch.no_grad()
def benchmark(model, inp):
    for i in range(100):
        model(inp)
    start = time.time()
    for i in range(200):
        model(inp)
    elapsed_ms = (time.time() - start) * 1000
    print("Latency: {:.2f}".format(elapsed_ms / 200))

# 對原始模型測速。
benchmark(model, dummy)

# 對最佳化後的模型測速。
benchmark(optimized_model, dummy)

擴充

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

直接傳入torch.jit.ScriptModule對象
從檔案載入torch.jit.save匯出的torch.jit.ScriptModule模型檔案。

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

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

後續步驟

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