PAI-Blade敏捷版支持TensorFlow和PyTorch模型在GPU硬件、端侧设备的INT8量化。本文主要介绍在GPU硬件上,如何使用PAI-Blade量化优化。

背景信息

量化是模型压缩常用手段之一,通过将原始的全精度32位浮点数分桶量化为位宽更小的定点整数,可以达到节省访存开销、提升指令计算吞率的双重目的(需要底层计算硬件的支持)。

TensorFlow模型量化

对TensorFlow模型进行基本优化的方法详情请参见优化TensorFlow模型。此外,您还可以开启量化功能,即使用PAI-Blade敏捷版优化时指定optimization_level='o2'。如果硬件支持INT8计算并存在量化性能收益,则PAI-Blade敏捷版会以默认方式进行量化。其中:
  • 如果不额外提供校正数据集(Calibration_Dataset),则PAI-Blade敏捷版执行在线INT8量化方案。
  • 如果希望获得进一步的加速效果,建议您提供一组用于离线计算量化参数的校正数据集,PAI-Blade敏捷版将自动开启离线INT8量化。
对于TensorFlow模型,校正数据集是一个包含若干组feed_dict的列表,示例如下。
# 准备校正数据集。
import numpy as np
calib_data = list()
for i in range(10):
    # feed dict中value必须全部为np.ndarray类型。
    feed_dict = {'input:0': np.ones((32, 224, 224, 3), dtype=np.float32)}
    calib_data.append(feed_dict)

您可以参考如下步骤对TensorFlow模型进行量化优化。

  1. 下载示例模型、测试数据及校正数据集。
    wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/tf_resnet50_v1.5/frozen.pb
    wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/tf_resnet50_v1.5/test_bc32.npy
    wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/tf_resnet50_v1.5/calib_data_test_bc32.npy
  2. 加载待优化模型和相应数据。
    import numpy as np
    import tensorflow as tf
    # 加载模型。
    graph_def = tf.GraphDef()
    with open('frozen.pb', 'rb') as f:
        graph_def.ParseFromString(f.read())
    # 加载测试数据。
    test_data = np.load('test_bc32.npy', allow_pickle=True, encoding='bytes',).item()
    # 加载校正数据集。
    calib_data = np.load('calib_data_test_bc32.npy', allow_pickle=True, encoding='bytes',)
  3. 执行离线INT8优化。
    import blade
    optimized_model, opt_spec, report = blade.optimize(
        model=graph_def,
        optimization_level='o2',
        device_type='gpu',
        test_data=test_data,
        calib_data=calib_data
    )
  4. 验证优化后的模型精度。

    量化优化完成后,您可以在完整的测试数据集上验证是否存在明显的精度损失。如果量化后的模型仍然满足精度要求,则无需继续执行后续尝试,反之可以尝试修改量化配置(quant_config),以缓解精度损失。关于量化配置的详细信息请参见表 1

    目前,TensorFlow模型在GPU上的量化配置仅支持weight_adjustment,启用该设置后,PAI-Blade敏捷版将自动调整模型参数以减小量化精度损失。启用该设置的示例代码如下所示。
    quant_config = {
        'weight_adjustment': 'true'  # 默认为'false'。
    }
    optimized_model, opt_spec, report = blade.optimize(
        model=graph_def,
        optimization_level='o2',
        device_type='gpu',
        test_data=test_data,
        calib_data=calib_data,
        config=blade.Config(quant_config=quant_config)
    )

PyTorch模型量化

与TensorFlow模型量化类似,您只需要使用PAI-Blade优化时指定optimization_level='o2'即可开启量化功能。但是与TensorFlow模型量化不同,PyTorch模型量化仅支持离线方式。因此,开启量化功能时,需要提供一组用于离线计算量化参数的校正数据集。

PyTorch模型的校正数据集是一个包含若干组输入数据的列表,示例如下。
# 准备校正数据集。
import numpy as np
calib_data = list()
for i in range(10):
    image = torch.ones(32, 3, 224, 224)
    calib_data.append(image)
您可以参考如下步骤对PyTorch模型进行量化优化。
  1. 下载示例模型、测试数据及校正数据集。
    wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/pt_resnet50_v1.5/traced_model.pt
    wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/pt_resnet50_v1.5/test_bc32.pth
    wget https://pai-blade.oss-cn-zhangjiakou.aliyuncs.com/test_public_model/bbs/pt_resnet50_v1.5/calib_data_test_bc32.pth
  2. 加载待优化模型和相应数据。
    import torch
    # 加载模型。
    pt_model = torch.jit.load('traced_model.pt')
    # 加载测试数据。
    test_data = torch.load('test_bc32.npy')
    # 加载校正数据集。
    calib_data = torch.load('calib_data_test_bc32.npy')
  3. 执行离线INT8优化。
    import blade
    optimized_model, opt_spec, report = blade.optimize(
        model=pt_model,
        optimization_level='o2',
        device_type='gpu',
        test_data=test_data,
        calib_data=calib_data
    )