更多 TVM 中文文档可访问 → Apache TVM 是一个端到端的深度学习编译框架,适用于 CPU、GPU 和各种机器学习加速芯片。 | Apache TVM 中文站
作者: Wuwei Lin
本文介绍如何用 TVM 自动量化(TVM 的一种量化方式)。有关 TVM 中量化的更多详细信息,参阅 此处。本教程将在 ImageNet 上导入一个 GluonCV 预训练模型到 Relay,量化 Relay 模型,然后执行推理。
准备数据集
以下演示如何为量化准备校准数据集,首先下载 ImageNet 的验证集,并对数据集进行预处理。
把校准数据集(可迭代对象)定义为 Python 中的生成器对象,本教程仅用几个样本进行校准。
导入模型
用 Relay MxNet 前端从 Gluon 模型集合(model zoo)中导入模型。
量化模型
量化过程要找到每一层的每个权重和中间特征图(feature map)张量的 scale。
对于权重而言,scales 是根据权重的值直接计算出来的。支持两种模式:power2 和 max。这两种模式都是先找到权重张量内的最大值。在 power2 模式下,最大值向下舍入为 2 的幂。如果权重和中间特征图的 scale 都是 2 的幂,则可以利用移位(bit shifting)进行乘法运算,这使得计算效率更高。在 max 模式下,最大值用作 scale。如果不进行四舍五入,在某些情况下 max 模式可能具有更好的精度。当 scale 不是 2 的幂时,将使用定点乘法。
中间特征图可以通过数据感知量化来找到 scale。数据感知量化将校准数据集作为输入参数,通过最小化量化前后激活分布之间的 KL 散度来计算 scales。或者也可以用预定义的全局 scales,这样可以节省校准时间,但会影响准确性。
运行推理
创建一个 Relay VM 来构建和执行模型。
输出结果:
脚本总运行时长: (1 分 22.338 秒)