AICompiler是集成在PAI-Blade中的AI编译优化组件,包含Static Shape和Dynamic Shape编译框架。通常您无需提供额外配置,AICompiler即可在通用透明的情况下帮助您提高推理性能。本文介绍如何使用AICompiler对TensorFlow和PyTorch模型进行编译优化。

背景信息

随着AI模型结构的快速演化,底层计算硬件的层出不穷,用户使用习惯的推陈出新,使得单纯基于手工优化解决AI模型的性能和效率问题越来越困难。为了解决这些问题,AI编译优化技术已经成为一个被广泛关注的技术方向。

传统编译器是以高层语言作为输入,避免用户直接写机器码。深度学习编译器的作用与其相仿,输入是比较灵活的、具备较高抽象度的计算图,输出包括CPU或GPU等硬件平台上的底层机器码及执行引擎。AI编译器的目标是针对AI计算任务,以通用编译器的方式完成性能优化。让用户可以专注于上层模型开发,降低用户手工优化性能的人力开发成本,进一步压榨硬件性能空间。

在过去两年多时间里,PAI团队在AI编译优化技术方向投入了比较专注的资源精力,AICompiler已经作为优化组件之一集成到PAI-Blade之中,从而帮助用户以通用透明的方式完成推理模型的优化和部署。

目前AICompiler主要包含Static Shape和Dynamic Shape两套编译框架。Dynamic Shape可以支持所有任务类型,包括计算图Shape变化范围较大的推理任务。Static Shape主要适配计算图为静态Shape的任务,或Shape变化范围较小的任务,以获得理论上更加极致的性能。Static Shape和Dynamic Shape的关系如下表所示。
编译框架 TensorFlow PyTorch 适配范围
Static Shape 支持 暂不支持 适配静态Shape或Shape变化范围较小的任务,理论上可以获得极致的性能收益。
Dynamic Shape 支持 支持 适配所有的任务类型。
默认情况下,PAI-Blade会自动分析用户的模型是否适合Dynamic Shape的编译优化,用户无需提供额外的输入信息。对于静态Shape类任务或Shape变化范围较小的任务,用户可以提供额外的输入选项,此时PAI-Blade会尝试通过Static Shape模式编译以获得更优的性能。下文详细介绍几个应用例子供您参考。

通过Dynamic Shape模式编译TensorFlow模型

以一个开源ASR模型为例,演示如何通过Dynamic Shape模式编译TensorFlow模型:
  1. 下载模型和测试数据。
    # 下载示例模型、测试数据。
    wget https://pai-blade.cn-hangzhou.oss.aliyun-inc.com/test_public_model/bbs/tf_aicompiler_demo/frozen.pb
    wget https://pai-blade.cn-hangzhou.oss.aliyun-inc.com/test_public_model/bbs/tf_aicompiler_demo/test_bc4.npy
  2. 加载模型、指定数据,并调用blade.optimize,无需任何额外配置。
    import numpy as np
    import tensorflow as tf
    import blade
    
    # 加载模型和测试数据。
    graph_def = tf.GraphDef()
    with open('./frozen.pb', 'rb') as f:
        graph_def.ParseFromString(f.read())
    test_data = np.load('test_bc4.npy', allow_pickle=True, encoding='bytes',).item()
    
    # 开始优化。
    optimized_model, opt_spec, report = blade.optimize(
        graph_def,  # The original model, here is a TF GraphDef.
        'o1',  # Optimization level o1 or o2.
        device_type='gpu',  # Target device to run the optimized model.
        config=blade.Config(),
        inputs=['encoder_memory_placeholder', 'encoder_memory_length_placeholder'],
        outputs=['score', 'seq_id'],
        test_data=[test_data]
        #verbose=True
    )
    
    # 保存优化结果。
    tf.train.write_graph(optimized_model, "./", "optimized.pb", as_text=False)
    print("Report: {}".format(report))
  3. 优化完成后,查看优化报告(blade.optimize返回的report)中AICompiler生效之后的性能收益。本例中以T4卡为例,AICompiler可以帮您在通用透明的情况下获得2.23倍的性能收益,report中的字段含义请参见优化报告
        {
          "name": "TfAicompilerGpu",
          "status": "effective",
          "speedup": "2.23",
          "pre_run": "120.54 ms",
          "post_run": "53.99 ms"
        }

通过Static Shape模式编译TensorFlow模型

如果模型部署后的实际Shape变化范围较小或完全是Static Shape,您可以通过配置config选项,指定PAI-Blade以Static Shape模式进行编译。相比Dynamic Shape模式编译,您只需修改blade.optimize函数的config入参即可实现Static Shape编译,示例代码如下。
optimized_model, opt_spec, report = blade.optimize(
    graph_def,  # The original model, here is a TF GraphDef.
    'o1',  # Optimization level o1 or o2.
    device_type='gpu',  # Target device to run the optimized model.
    # Provide an additional config here in order to try Static Shape Compilation:
    config=blade.Config(enable_static_shape_compilation_opt = True), 
    inputs=['encoder_memory_placeholder', 'encoder_memory_length_placeholder'],
    outputs=['score', 'seq_id'],
    test_data=[test_data]
    #verbose=True
)
关于config参数的高级配置,详情请参见表 1
优化完成后,可以查看优化报告(blade.optimize返回的report)中的性能收益。本例中的性能收益进一步提升至2.35倍,report中的字段含义请参见优化报告
    {
      "name": "TfAicompilerGpu",
      "status": "effective",
      "speedup": "2.35",
      "pre_run": "114.91 ms",
      "post_run": "48.86 ms"
    }

通过Dynamic Shape模式编译PyTorch模型

以一个开源ASR模型为例,演示如何通过Dynamic Shape模式编译PyTorch模型:
  1. 下载模型。
    # PyTorch 1.6.0
    # Python3.6
    wget https://pai-blade.cn-hangzhou.oss.aliyun-inc.com/test_public_model/bbs/pt_aicompiler_demo/orig_decoder_v2.pt
  2. 加载模型、指定数据,并调用blade.optimize,无需任何额外配置。
    import os
    import time
    import torch
    # To use blade, just import it.
    import blade
    
    # 加载模型。
    pt_file = 'orig_decoder_v2.pt'
    batch = 8
    model = torch.jit.load(pt_file)
    
    # 准备测试数据。
    def get_test_data(batch_size=1):
        decoder_input_t = torch.LongTensor([1] * batch_size).cuda()
        decoder_hidden_t = torch.rand(batch_size, 1, 256).cuda()
        decoder_hidden_t = decoder_hidden_t * 1.0
        decoder_hidden_t = torch.tanh(decoder_hidden_t)
        output_highfeature_t = torch.rand(batch_size, 448, 4, 50).cuda()
        attention_sum_t = torch.rand(batch_size, 1, 4, 50).cuda()
        decoder_attention_t = torch.rand(batch_size, 1, 4, 50).cuda()
        et_mask = torch.rand(batch_size, 4, 50).cuda()
    
        return (decoder_input_t, decoder_hidden_t, output_highfeature_t, attention_sum_t, decoder_attention_t, et_mask)
    
    dummy = get_test_data(batch)
    
    # 开始优化。
    optimized_model, opt_spec, report = blade.optimize(
        model,  # The original model, here is a torch scrip model.
        'o1',  # Optimization level o1 or o2.
        device_type='gpu',  # Target device to run the optimized model.
        test_data=[dummy],  # For PyTorch, input data is list of tupoles.
        config=blade.Config()
    )
    
    print("spec: {}".format(opt_spec))
    print("report: {}".format(report))
    
    # 导出优化结果。
    torch.jit.save(optimized_model, 'optimized_decoder.pt')
  3. 优化完成后,查看优化报告(blade.optimize返回的report)中AICompiler生效之后的性能收益。本例中以T4卡为例,AICompiler可以帮您在通用透明的情况下获得2.45倍的性能收益,report中的字段含义请参见优化报告
      "optimizations": [
        {
          "name": "PyTorchMlir",
          "status": "effective",
          "speedup": "2.45",
          "pre_run": "1.99 ms",
          "post_run": "0.81 ms"
        }
      ],