如何在GPU硬件上使用PAI-Blade量化优化 - 人工智能平台 PAI

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模型进行量化优化。

下载示例模型、测试数据及校正数据集。

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

加载待优化模型和相应数据。

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',)

执行离线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
)

验证优化后的模型精度。
量化优化完成后，您可以在完整的测试数据集上验证是否存在明显的精度损失。如果量化后的模型仍然满足精度要求，则无需继续执行后续尝试，反之可以尝试修改量化配置（quant_config），以缓解精度损失。关于量化配置的详细信息请参见blade.Config。
目前，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模型进行量化优化。

下载示例模型、测试数据及校正数据集。

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

加载待优化模型和相应数据。

import torch
# 加载模型。
pt_model = torch.jit.load('traced_model.pt')
# 加载测试数据。
test_data = torch.load('test_bc32.pth')
# 加载校正数据集。
calib_data = torch.load('calib_data_test_bc32.pth')

执行离线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
)