文档

优化PyTorch模型

更新时间:

PAI-Blade提供了丰富的模型优化方法,您只需要在本地环境中安装Wheel包,即可通过调用Python API的方式进行模型优化。本文介绍如何使用PAI-Blade优化PyTorch模型,所有实验结果均在NVidia T4卡上测得。

前提条件

  • 已安装PyTorch及PAI-Blade的Wheel包,详情请参见安装Blade

  • 已有训练完成的PyTorch模型,本文使用一个公开的ResNet50模型。

优化PyTorch模型

  1. 导入PAI-Blade和其他依赖库。

    import os
    import time
    import torch
    import torchvision.models as models
    import blade
  2. 从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() # 构造测试数据。
  3. 调用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.
  4. 打印优化报告。

    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"         // 加速比。
      },
      // ......
    }
  5. 对比优化前后的性能。

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

  • 本页导读 (1)
文档反馈