优化PyTorch模型_人工智能平台 PAI(PAI)-阿里云帮助中心

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模型推理。