優(yōu)化PyTorch模型

1. 導(dǎo)入PAI-Blade和其他依賴庫(kù)。

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

2. 從torchvision加載ResNet50模型。由于PAI-Blade僅支持ScriptModule，因此需要轉(zhuǎn)換模型格式。

   model = models.resnet50().float().cuda()  # 準(zhǔn)備模型。
   model = torch.jit.script(model).eval()    # 轉(zhuǎn)換成ScriptModule。
   dummy = torch.rand(1, 3, 224, 224).cuda() # 構(gòu)造測(cè)試數(shù)據(jù)。

3. 調(diào)用blade.optimize函數(shù)進(jìn)行模型優(yōu)化，詳細(xì)的參數(shù)解釋請(qǐng)參見Python接口文檔。模型優(yōu)化的代碼示例如下。如果您在優(yōu)化中遇到問(wèn)題，可以加入Blade客戶群，咨詢相關(guān)人員，詳情請(qǐng)參見獲取Token。

   optimized_model, opt_spec, report = blade.optimize(
       model,                 # 待優(yōu)化的模型。
       'o1',                  # 優(yōu)化級(jí)別，o1或o2。
       device_type='gpu',     # 目標(biāo)設(shè)備，gpu/cpu。
       test_data=[(dummy,)],  # PyTorch的輸入數(shù)據(jù)是List of tuple of tensor。
   )

   blade.optimize函數(shù)返回的三個(gè)對(duì)象，分別如下所示：

   • optimized_model：優(yōu)化完成的模型，此處為torch.jit.ScriptModule。

   • opt_spec：包含復(fù)現(xiàn)優(yōu)化結(jié)果需要的配置信息、環(huán)境變量及資源文件等，通過(guò)with語(yǔ)句可以使其生效。

   • report：優(yōu)化報(bào)告，可以直接打印。關(guān)于報(bào)告中的參數(shù)解釋，詳情請(qǐng)參見優(yōu)化報(bào)告。

   優(yōu)化過(guò)程中，您可以看到如下類似的優(yōu)化進(jìn)度。

   [Progress] 5%, phase: user_test_data_validation.
   [Progress] 10%, phase: test_data_deduction.
   [Progress] 15%, phase: CombinedSwitch_4.
   [Progress] 95%, phase: model_collecting.

4. 打印優(yōu)化報(bào)告。

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

   在優(yōu)化報(bào)告中可以看到主要的效果源于哪些優(yōu)化項(xiàng)，如下所示。

   Report: {
     // ......
     "optimizations": [
       {
         "name": "PtTrtPassFp32",
         "status": "effective",
         "speedup": "1.50",     // 加速比。
         "pre_run": "5.29 ms",  // 加速前延遲。
         "post_run": "3.54 ms"  // 加速后延遲。
       }
     ],
     // 端到端優(yōu)化結(jié)果。
     "overall": {
       "baseline": "5.30 ms",    // 原始模型延遲。
       "optimized": "3.59 ms",   // 優(yōu)化后模型延遲。
       "speedup": "1.48"         // 加速比。
     },
     // ......
   }

5. 對(duì)比優(yōu)化前后的性能。

   @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))
   
   # 對(duì)原始模型測(cè)速。
   benchmark(model, dummy)
   
   # 對(duì)優(yōu)化后的模型測(cè)速。
   benchmark(optimized_model, dummy)

擴(kuò)展

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

后續(xù)步驟

經(jīng)過(guò)PAI-Blade優(yōu)化的模型，您可以通過(guò)Python直接執(zhí)行或部署為EAS服務(wù)。此外，PAI-Blade也提供了C++ SDK，以便您將優(yōu)化后的模型集成到自己的應(yīng)用中，詳情請(qǐng)參見使用SDK部署PyTorch模型推理。