常规推理优化普遍针对输入为Static Shape的模型,如果实际推理的模型Shape发生变化,推理优化效果就可能失效。在实际生产中,输入为Dynamic Shape的模型越来越多,因此对不同输入Shape的推理过程具有强烈的优化需求。本文介绍如何使用Blade优化输入为Dynamic Shape的模型。

使用限制

本文使用的环境需要满足以下版本要求:
  • 系统环境:Linux系统中使用Python 3.6及其以上版本。
  • 框架:PyTorch 1.7.1。
  • 设备及后端:NVIDIA T4、CUDA 11.0。
  • 推理优化工具:Blade 3.17.0及其以上版本。

操作流程

使用Blade优化输入为Dynamic Shape的ResNet50流程如下:
  1. 步骤一:准备工作

    构建测试数据和模型,本文使用torchvision中标准的ResNet50模型。

  2. 步骤二:配置用于优化的config

    根据Dynamic Shape的范围配置Blade config。

  3. 步骤三:调用Blade优化模型

    调用blade.optimize接口优化模型,并保存优化后的模型。

  4. 步骤四:验证性能与正确性

    对优化前后的推理速度及推理结果进行测试,从而验证优化报告中信息的正确性。

  5. 步骤五:加载运行优化后的模型

    集成Blade SDK,加载优化后的模型进行推理。

步骤一:准备工作

  1. 下载模型预训练参数与测试数据。
    预训练参数选自torchvision,为了加速下载过程,已将其存储至OSS中。测试数据随机选自ImageNet-1k验证集,预处理操作已完成,您可以下载后直接使用。
    !wget http://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/share/dynamic_ranges_pratice/resnet50-19c8e357.pth -O resnet50-19c8e357.pth
    !wget http://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/share/dynamic_ranges_pratice/imagenet_val_example.pt -O imagenet_val_example.pt
  2. 定义模型、加载模型参数和测试数据,并生成TorchScript。
    import torch
    import torchvision
    
    # 构建Resnet50。
    model = torchvision.models.resnet50().eval().cuda()
    # 加载预训练参数。
    ckpt = torch.load('resnet50-19c8e357.pth')
    model.load_state_dict(ckpt)
    # 加载测试数据。
    example_input = torch.load('imagenet_val_example.pt').cuda()
    # 生成TorchScript。
    traced_model = torch.jit.trace(model, example_input).cuda().eval()

步骤二:配置用于优化的config

根据Dynamic Shape的范围配置Blade config,Blade支持任意维度的动态范围。本文以Batch维度演示config的配置。

  1. 定义Dynamic Shape的范围。
    一组有效的动态范围,需要包括以下三个字段:
    • min:表示dynamic shape的下界。
    • max:表示dynamic shape的上界。
    • opts:表示需要特别优化的Shape,可以设置多个。通常优化后的模型在这些Shape上的推理加速比更高。
    上述三个字段需要符合以下规则:
    • minmaxopts中的每组Shape的长度相等,且等于网络的输入数量。
    • minmaxopts中的每组Shape对应位置的数值需要满足min_num <= opt_num <= max_num
    例如构建如下Dynamic Shape的范围。
    shapes = {
        "min": [[1, 3, 224, 224]],
        "max": [[10, 3, 224, 224]],
        "opts": [
            [[5, 3, 224, 224]],
            [[8, 3, 224, 224]],
        ]
    }
    此外,Blade支持设置多个动态范围。如果Dynamic Shape的上界和下界范围过大,可能会导致优化后的模型加速不明显,您可以将一个大的范围拆分为多个小范围,通常能够带来更好的加速效果。关于如何设置多个动态范围,请参见下文的附录:设置多个动态范围
  2. 通过定义好的Dynamic Shape范围构建Blade config。
    import blade
    import blade.torch as blade_torch
    
    # Blade Torch相关config,用于设置Dynamic Shapes。
    blade_torch_cfg = blade_torch.Config()
    blade_torch_cfg.dynamic_tuning_shapes = shapes
    
    # Blade相关config,用于关闭FP16的精度检查,以获得最好的加速效果。
    gpu_config = {
        "disable_fp16_accuracy_check": True,
    }
    blade_config = blade.Config(
        gpu_config=gpu_config
    )

步骤三:调用Blade优化模型

  1. 调用blade.optimize对模型进行优化,示例代码如下。关于该接口的详细描述,请参见Python接口文档
    with blade_torch_cfg:
        optimized_model, _, report = blade.optimize(
            traced_model,          # 模型路径。
            'o1',                  # o1无损优化。
            config=blade_config,
            device_type='gpu',     # 面向GPU设备优化,
            test_data=[(example_input,)]  # 测试数据。
        )
    优化模型时,您需要注意以下事宜:
    • blade.optimize的第一个返回值为优化后的模型,其数据类型与输入的模型相同。在这个示例中,输入的是TorchScript,返回的是优化后的TorchScript。
    • 您需要确保输入的test_data在定义的Dynamic Shape范围内。
  2. 优化完成后,打印优化报告。
    print("Report: {}".format(report))
    打印的优化报告类似如下输出。
    Report: {
      "software_context": [
        {
          "software": "pytorch",
          "version": "1.7.1+cu110"
        },
        {
          "software": "cuda",
          "version": "11.0.0"
        }
      ],
      "hardware_context": {
        "device_type": "gpu",
        "microarchitecture": "T4"
      },
      "user_config": "",
      "diagnosis": {
        "model": "unnamed.pt",
        "test_data_source": "user provided",
        "shape_variation": "undefined",
        "message": "Unable to deduce model inputs information (data type, shape, value range, etc.)",
        "test_data_info": "0 shape: (1, 3, 224, 224) data type: float32"
      },
      "optimizations": [
        {
          "name": "PtTrtPassFp16",
          "status": "effective",
          "speedup": "4.06",
          "pre_run": "6.55 ms",
          "post_run": "1.61 ms"
        }
      ],
      "overall": {
        "baseline": "6.54 ms",
        "optimized": "1.61 ms",
        "speedup": "4.06"
      },
      "model_info": {
        "input_format": "torch_script"
      },
      "compatibility_list": [
        {
          "device_type": "gpu",
          "microarchitecture": "T4"
        }
      ],
      "model_sdk": {}
    }
    从优化报告可以看出本示例的优化中,PtTrtPassFp16优化项生效,带来了约4.06倍左右的加速,将模型在测试数据上的推理耗时从6.55 ms下降到了1.61 ms。上述优化结果仅为本示例的测试结果,您的优化效果以实际为准。关于优化报告的字段详情请参见优化报告
  3. 调用PyTorch的相关函数保存并加载优化后的TorchScript模型。
    file_name = "resnet50_opt.pt"
    # 将优化后的模型保存到本地。
    torch.jit.save(optimized_model, file_name)
    # 从硬盘中加载优化后的模型。
    optimized_model = torch.jit.load(file_name)

步骤四:验证性能与正确性

优化完成后,通过Python脚本对优化报告的信息进行验证。

  1. 定义benchmark方法,对模型进行10次预热,然后运行100次,最终取平均的推理时间作为推理速度。
    import time
    
    @torch.no_grad()
    def benchmark(model, test_data):
        # 切换模型至验证模式。
        model = model.eval()
        
        # 预热。
        for i in range(0, 10):
            model(test_data)
            
        # 开始计时运行。
        num_runs = 100
        start = time.time()
        for i in range(0, num_runs):
            model(test_data)
        torch.cuda.synchronize()
        elapsed = time.time() - start
        rt_ms = elapsed / num_runs * 1000.0
        
        # 打印结果。
        print("{:.2f} ms.".format(rt_ms))
        return rt_ms
  2. 定义一系列不同Shape的测试数据。
    dummy_inputs = []
    batch_num = [1, 3, 5, 7, 9]
    for n in batch_num:
        dummy_inputs.append(torch.randn(n, 3, 224, 224).cuda())
  3. 遍历每组测试数据,分别调用benchmark方法对优化前与优化后的模型进行测试,并打印结果打印。
    for inp in dummy_inputs:
        print(f'--------------test with shape {list(inp.shape)}--------------')
        print("  Origin model inference cost:     ", end='')
        origin_rt = benchmark(traced_model, inp)
        print("  Optimized model inference cost:  ", end='')
        opt_rt = benchmark(optimized_model, inp)
        speedup = origin_rt / opt_rt
        print('  Speed up: {:.2f}'.format(speedup))
        print('')
    系统返回如下类似结果。
    --------------test with shape [1, 3, 224, 224]--------------
      Origin model inference cost:     6.54 ms.
      Optimized model inference cost:  1.66 ms.
      Speed up: 3.94
    
    --------------test with shape [3, 3, 224, 224]--------------
      Origin model inference cost:     10.79 ms.
      Optimized model inference cost:  2.40 ms.
      Speed up: 4.49
    
    --------------test with shape [5, 3, 224, 224]--------------
      Origin model inference cost:     16.27 ms.
      Optimized model inference cost:  3.25 ms.
      Speed up: 5.01
    
    --------------test with shape [7, 3, 224, 224]--------------
      Origin model inference cost:     22.62 ms.
      Optimized model inference cost:  4.39 ms.
      Speed up: 5.16
    
    --------------test with shape [9, 3, 224, 224]--------------
      Origin model inference cost:     28.83 ms.
      Optimized model inference cost:  5.25 ms.
      Speed up: 5.49
    从结果可以看出对于不同Shape的测试数据,优化后模型的推理速度是原始模型的3.94~5.49倍。上述优化结果仅为本示例的测试结果,您的优化效果以实际为准。
  4. 使用准备工作阶段准备的真实测试数据example_input,验证优化模型的正确性。
    origin_output = traced_model(example_input)
    _, pred = origin_output.topk(1, 1, True, True)
    print("origin model output: {}".format(pred))
    opt_output = optimized_model(example_input)
    _, pred = origin_output.topk(1, 1, True, True)
    print("optimized model output: {}".format(pred))
    系统返回如下类似结果。
    origin model output: tensor([[834]], device='cuda:0')
    optimized model output: tensor([[834]], device='cuda:0')
    从上述结果可以看出优化前后模型对于测试数据example_input的预测均为第834类。

步骤五:加载运行优化后的模型

完成验证后,您需要对模型进行部署,Blade敏捷版提供了Python和C++两种运行时SDK供您集成。关于C++的SDK使用方法请参见使用SDK部署TensorFlow模型推理,下文主要介绍如何使用Python SDK部署模型。

  1. 可选:在试用阶段,您可以设置如下的环境变量,防止因为鉴权失败而程序退出。
    export BLADE_AUTH_USE_COUNTING=1
  2. 获取鉴权。
    export BLADE_REGION=<region>
    export BLADE_TOKEN=<token>
    您需要根据实际情况替换以下参数:
    • <region>:Blade敏捷版支持的地域,需要加入Blade用户群获取该信息,用户群的二维码详情请参见获取Token
    • <token>:鉴权Token,需要加入Blade用户群获取该信息,用户群的二维码详情请参见获取Token
  3. 加载运行优化后的模型。
    除了增加一行import blade.runtime.torch,您无需为Blade的接入编写额外代码,即原有的推理代码无需任何改动。
    import torch
    import blade.runtime.torch
    # <your_optimized_model_path>替换为优化后的模型路径。
    opt_model_dir = <your_optimized_model_path>
    # <your_infer_data>替换为用于推理的数据。
    infer_data = <your_infer_data>
    
    model = torch.jit.load(opt_model_dir)
    output = model(infer_data)

附录:设置多个动态范围

如果Dynamic Shape的上界和下界范围过大,可能会导致优化后的模型加速不明显,您可以将一个大的范围拆分为多个小范围,通常能够带来更好的加速效果。例如设置如下Dynamic Shape。
shapes1 = {
    "min": [[1, 3, 224, 224]],
    "max": [[5, 3, 224, 224]],
    "opts": [
        [[5, 3, 224, 224]],
    ]
}
shapes2 = {
    "min": [[5, 3, 224, 224]],
    "max": [[10, 3, 224, 224]],
    "opts": [
        [[8, 3, 224, 224]],
    ]
}
shapes = [shapes1, shapes2]
您可以使用该shapes配置上述提及的优化config,详情请参见步骤二:配置用于优化的config