PAI-Blade提供了豐富的模型優化方法，您只需要在本地環境中安裝Wheel包，即可通過調用Python API的方式進行模型優化。本文介紹如何使用PAI-Blade優化PyTorch模型，所有實驗結果均在NVidia T4卡上測得。
前提條件
已安裝PyTorch及PAI-Blade的Wheel包，詳情請參見安裝Blade。
已有訓練完成的PyTorch模型，本文使用一個公開的ResNet50模型。
優化PyTorch模型
導入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模型推理。