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量化。
# 准备校正数据集。
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模型进行量化优化。
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模型进行量化优化。