全部产品
Search
文档中心

人工智能平台 PAI:优化TensorFlow模型

更新时间:Mar 04, 2024

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

前提条件

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

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

优化TensorFlow模型

本文以一个公开的ResNet50模型为例,演示如何优化TensorFlow模型。您也可以对自己的TensorFlow模型进行优化。

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

    import os
    import numpy as np
    import tensorflow.compat.v1 as tf
    import blade
  2. 编写一个简单的函数,用于下载待优化的模型和测试数据。

    虽然PAI-Blade支持没有测试数据的优化,即零输入优化,但是基于真实输入数据的优化结果会更加准确有效。因此,建议您提供测试数据。下载待优化的模型和测试数据的函数示例如下所示。

    def _wget_demo_tgz():
        # 下载一个公开的resnet50模型。
        url = 'http://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/demo/mask_rcnn_resnet50_atrous_coco_2018_01_28.tar.gz'
        local_tgz = os.path.basename(url)
        local_dir = local_tgz.split('.')[0]
        if not os.path.exists(local_dir):
            blade.util.wget_url(url, local_tgz)
            blade.util.unpack(local_tgz)
        model_path = os.path.abspath(os.path.join(local_dir, "frozen_inference_graph.pb"))
        graph_def = tf.GraphDef()
        with open(model_path, 'rb') as f:
            graph_def.ParseFromString(f.read())
        # 以随机数作为测试数据。
        test_data = np.random.rand(1, 800,1000, 3)
        return graph_def, {'image_tensor:0': test_data}
    
    graph_def, test_data = _wget_demo_tgz()
  3. 调用blade.optimize函数进行模型优化,详细的参数解释请参见Python接口文档。模型优化的代码示例如下。

    input_nodes=['image_tensor']
    output_nodes = ['detection_boxes', 'detection_scores', 'detection_classes', 'num_detections', 'detection_masks']
    
    optimized_model, opt_spec, report = blade.optimize(
        graph_def,                 # 待优化的模型,此处是tf.GraphDef, 也可以配置为SavedModel的路径。
        'o1',                      # 优化级别,o1或o2。
        device_type='gpu',         # 目标设备,gpu/cpu/edge。
        inputs=input_nodes,        # 输入节点,也可以不提供,PAI-Blade会自行推断。
        outputs=output_nodes,      # 输出节点。
        test_data=[test_data]      # 测试数据。
    )

    blade.optimize函数返回的三个对象,分别如下所示:

    • optimized_model:优化完成的模型,本文为tf.GraphDef

    • opt_spec:包含复现优化结果需要的配置信息、环境变量及资源文件等,通过with语句可以使其生效。

    • report:优化报告,可以直接打印。关于报告中的参数解释,详情请参见优化报告

    优化过程中,您可以看到如下类似的优化进度。

    [Progress] 5%, phase: user_test_data_validation.
    [Progress] 10%, phase: test_data_deduction.
    [Progress] 15%, phase: CombinedSwitch_1.
    [Progress] 24%, phase: TfStripUnusedNodes_22.
    [Progress] 33%, phase: TfStripDebugOps_23.
    [Progress] 42%, phase: TfFoldConstants_24.
    [Progress] 51%, phase: CombinedSequence_7.
    [Progress] 59%, phase: TfCudnnrnnBilstm_25.
    [Progress] 68%, phase: TfFoldBatchNorms_26.
    [Progress] 77%, phase: TfNonMaxSuppressionOpt_27.
    [Progress] 86%, phase: CombinedSwitch_20.
    [Progress] 95%, phase: model_collecting.
    [Progress] 100%, Finished!
  4. 打印优化报告。

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

    在优化报告中可以看到主要的效果源于哪些优化项,如下所示。

    Report: {
      // ......
      "optimizations": [
        // ......
        {
          "name": "TfNonMaxSuppressionOpt",
          "status": "effective",
          "speedup": "1.58",        // 加速比。
          "pre_run": "522.74 ms",   // 优化前延迟。
          "post_run": "331.45 ms"   // 优化后延迟。
        },
        {
          "name": "TfAutoMixedPrecisionGpu",
          "status": "effective",
          "speedup": "2.43",
          "pre_run": "333.30 ms",
          "post_run": "136.97 ms"
        }
        // ......
      ],
      // 端到端优化结果。
      "overall": {
        "baseline": "505.91 ms",    // 原始模型延迟。
        "optimized": "136.83 ms",   // 优化后模型延迟。
        "speedup": "3.70"           // 加速比。
      },
      // ......
    }
  5. 对比优化前后的性能。

    import time
    
    def benchmark(model):
        tf.reset_default_graph()
        with tf.Session() as sess:
            sess.graph.as_default()
            tf.import_graph_def(model, name="")
            # Warmup!
            for i in range(0, 1000):
                sess.run(['image_tensor:0'], test_data)
            # Benchmark!
            num_runs = 1000
            start = time.time()
            for i in range(0, num_runs):
                sess.run(['image_tensor:0'], test_data)
            elapsed = time.time() - start
            rt_ms = elapsed / num_runs * 1000.0
            # Show the result!
            print("Latency of model: {:.2f} ms.".format(rt_ms))
    
    # 对原模型测速。
    benchmark(graph_def)
    
    # 对优化后的模型测速。
    with opt_spec:
        benchmark(optimized_model)

    性能实测结果如下,该结果与优化报告中的值基本一致。

    Latency of model: 530.26 ms.
    Latency of model: 148.40 ms.

扩展

blade.optimize函数的model参数支持多种形式的模型输入。对于TensorFlow模型,支持以下三种方式传入模型:

  • 直接传入tf.GraphDef对象

  • 从文件加载PB或PBTXT格式的Frozen PB

  • 从指定路径导入SavedModel

在本文示例中,为blade.optimize函数传入了内存中的tf.GraphDef对象。另外两种方式可以参考如下代码:

  • 传入Fronzen PB文件

    optimized_model, opt_spec, report = blade.optimize(
        './path/to/frozen_pb.pb',  # 也可以为.pbtxt格式。
        'o1',
        device_type='gpu',
    )
  • 传入SavedModel路径

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

后续步骤

经过PAI-Blade优化的模型,您可以通过Python直接执行或部署为EAS服务。此外,PAI-Blade也提供了C++ SDK,以便您将优化后的模型集成到自己的应用中,详情请参见使用SDK部署TensorFlow模型推理