探索 TVM 进行量化方法

 探索 TVM 进行量化方法

 

 Relay框架

如上图所示，有两种不同的并行工作正在进行中

• 自动整数量化 - 采用 FP32 框架图，在 Relay 中自动转换为 Int8。
• 接受预量化整数模型 - 这种方法接受预量化模型，引入称为 QNN 的Relay方言，生成 Int8 Relay图。

关于 Relay Automatic FP16 Downcasting 的讨论很少。还没有任何 RFC。 正在对此进行探索/原型设计，计划提出 RFC。

Relay优化

• 与目标无关的Relay pass- TVM 社区不断添加这些 pass。例子是fuse常量，公共子表达式消除等。
• 依赖于目标的Relay pass- 这些 pass转换Relay图，针对目标对进行优化。一个例子是 Legalize，或 AlterOpLayout 变换，改变卷积/密集层的布局。TVM 社区正在努力改进基础架构，实现此类转换，添加特定于目标的布局转换。一些基础架构工作，良好整体设计的先决条件。

Relay到硬件

有了优化的Relay图，就需要编写优化的调度。像 FP32 一样，必须只专注于昂贵的算子，如 conv2d、dense 等。有分散的努力，一些在不同后端工作的开发人员（不一定是 Int8），TVM 社区正在努力统一。

• Intel x86 - 近期 Int8 探索仅限于 Skylake 和 Cascade Lake
• ARM - 目前正在为 FP32 进行一些 NHWC 工作。计划是在 FP32 工作完成后，将这项工作扩展到 Int8。
• 英伟达—— 

基于搜索的自动量化

背景

在 tvm 中实现了一个量化工作流程，从一些现有的量化框架中，选择采用注释-校准-实现3阶段设计：

 

• Annotation：annotation pass根据每个算子的rewrite函数，重写graph，插入模拟的量化算子。模拟量化算子模拟，从浮点数到整数量化的舍入误差和饱和误差，
• 校准：校准通道将调整模拟量化算子的阈值，减少精度下降。
• 实现：实现过程将实际使用float32计算的模拟图，转换为真正的低精度整数图。

在开发过程中，存在一些缺点：

 

• 在注释中，每个注释作为张INPUT/ WEIGHT/ACTIVATION种不同的量化战略，这是种特殊算子的，需要不同的组合发生在不同的模式。这个 make annotation 有很多手动规则，变得很难维护。
• 模拟图没有总比例和数据类型信息。将尺度推理和数据类型选择，推迟到实现，这使得这部分的逻辑很难理解。此外，缺乏这些信息，无法在模拟过程中，捕获溢出错误。
• 尝试将量化模型，部署到不同硬件时，面临硬件差异。有两种解决方案： 1.注解时检查目标，逻辑比较复杂；2.添加一个新的partition pass，首先决定量化拓扑，每个硬件都需要实现一个定制的partition pass。

提出了一个新的量化框架，在循环中，引入了硬件和学习方法。已经进行了多项改进以，解决之前的问题：

• 在每条边上插入 SimQ (simulated_quantize) 算子，不是通过手动注释规则。让学习算法在每条边上，发现最佳量化策略，不是通过标记。
• 将in_scale, in_dtype，添加out_dtype到 SimQ 的定义中。在模拟期间，执行比例推理和数据类型分配。在 SimQ 中，模拟溢出错误。
• 提出Hardware抽象，描述硬件属性和算子约束。通过这种声明方式，用户只需Hardware，为不同的硬件定义不同的对象，无需了解量化逻辑。

工作流程概述

工作流程

 给定目标硬件的模型和描述，系统将生成一组，位的选择空间和Topology量化的空间。这里的Topology意思，考虑到硬件和算子约束，哪些节点/边将被量化，这将在后面讨论。

然后搜索循环开始：学习算法将从选择空间中，选择一组参数——每条边上的位数。阈值可以通过从小型校准数据集，收集的统计数据估计。结合拓扑、位和阈值，可以生成模拟图，在校准数据集（大约 128 个样本）上，对其进行评估。输出/精度作为反馈，学习算法可以选择下一组位设置。

最后，通过搜索找到的最佳策略，将模拟模型实现真实的低精度整数模型。

规格：位、阈值、标度

将介绍几种重要性符号：位、阈值、比例。

一般而言，量化的目标，将浮点数（实数值）运行的图，转换为整数（定量值）运行的图，不会牺牲太多精度。给定一个具有真实值的张量，转换后的 quant 值的关系是什么？这是将在当前实现中遵循的规范：

 

 硬件说明

硬件描述，试图为在量化过程中，需要考虑的硬件属性，提供一个中心抽象。通过声明这些属性，可以避免在随后的量化步骤中，处理硬件特定条件。

目前可以指定每个算子的，输入数据类型和输出数据类型。

desc = Hardware()

desc['add'].append(OpDesc(in_dtypes=['int32', 'int32'], out_dtypes=['int32']))

desc['add'].append(OpDesc(in_dtypes=['float32', 'float32'], out_dtypes=['float32']))

desc['nn.conv2d'].append(OpDesc(in_dtypes=['int16', 'int16'], out_dtypes=['int32']))

desc['nn.conv2d'].append(OpDesc(in_dtypes=['int8', 'int8'], out_dtypes=['int32']))

desc['nn.global_avg_pool2d'].append(OpDesc(in_dtypes=['float32', 'float32'], out_dtypes=['float32']))

硬件信息在整个过程中已经使用了多次：

• 通过指定算子只支持浮点计算，系统将实现一个结束，需要放在算子之前。可以解决 VTA 流水线的一些问题，指定一些算子，在 VTA 核心上，使用整数指令运行，一些算子，在普通 cpu 上，使用浮点指令。
• 位选择空间由此产生。对于每条边，可以推理出使用的最大位，取决于数据类型约束。
• 在决定了每条边使用的位数后，根据硬件信息，选择合适的数据类型。

模拟

阈值估计

为了估计阈值，在校准数据集上运行模型，收集需要的统计信息。目前将保存中间算子的所有输出。为了从收集的输出中确定阈值，有几种策略：

• max_range：使用输出的最大值作为对应节点的阈值。
• power2_range：将最大值四舍五入到最接近的两个值的幂，作为阈值。
• kl_estimate：选择一个阈值，使实际输出和量化输出之间，KL 距离足够小。

目前，选择了这种power2_range方法，可以使用移位来代替乘法，在最终的量化模型中，提供更好的性能。虽然kl_estimate带来更好的准确度，但相当耗时，目前在搜索中使用不可行。

一个棘手的问题是，对于像加法这样的算子，只能在其算子的标度为 eqaul 时执行。首先统一其算子的规模。为了实现这一点，估计阈值将在模拟之前进行调整。threshold_rectify引入了一个命名转换和一个特定于算子的属性：

@register_fthreshold_rectify('add')

def threshold_rectify_for_add(in_bits, out_bits, in_tholds, out_tholds):

   # choose scale of the one with maximum threshold

   idx = np.argmax(in_tholds)

   unified_scale = in_tholds[idx] / (2**(in_bits[idx] - sign_bit))

   # adjust thresholds according to the unified scale

   ...

模拟量化

给定比特和阈值，可以尝试生成一个模型，模拟量化带来的误差。经过分析，可以发现误差来自几个方面： 1.舍入误差；2.饱和误差；3.溢出错误。

将simulated_quantize在每条边上，插入一个算子，试图模拟这些错误。定义如下：

def simulated_quantize(data, in_scale, out_scale, clip_min, clip_max, in_dtype, out_dtype):

    if in_dtype == 'float32' and out_dtype == 'float32':

        # no need to quantize

        return data

    # simulated overflow error    

    data = data / in_scale

    data = topi.cast(data, in_dtype)

    data = data * in_scale

    scaled_data = data / out_scale

    # simulate saturated error

    clipped_data = topi.clip(scaled_data, clip_min, clip_max)

    # simulate round error

    rounded_data = topi.cast(topi.round(scaled_data), out_dtype)

    out = rounded_data * out_scale

    return out

如何通过位和阈值，计算这些参数呢？out_scale、clip_min、clip_max 是非常严格的：

integer_range = 2**(bit - sign_bit)

out_scale = threshold / integer_range

clip_min = - (integer_range - 1)

clip_max =    integer_range - 1

对于in_scale、in_dtype、out_dtype，需要做额外推理。

尺度推理

可以在上面的模型中，发现in_scale，SimQ 的实际上，前一个算子输出的尺度，可以根据算子定义计算。为这样的属性，提供了一个注册函数：

@register_finfer_scale('nn.conv2d'):

def infer_scale_for_conv2d(in_scales):

    return in_scales[0] * in_scales[1]

数据类型分配

对于数据类型，将遍历算子，从硬件描述中，选择满足输入位和输出位要求的算子规范。

学习

有了上面描述的所有准备工作，量化问题转换为学习问题：希望从选择空间中，找到最佳设置，以实现模拟模型的最佳精度（或其它目标，如性能），可以使用每轮的输出（准确度）作为反馈。

对于这个学习问题，实现了random_search, simulated_anealing, 也是一个贪心算法。目前实验表明贪婪搜索是最可行的。

日志格式

搜索空间很大，搜索过程可能很长，最好有一个正式的日志格式，记录实验细节，实现可重复性和可交换性。选择json格式，详细信息如下：

• version : 日志格式版本。
• 策略：量化策略。
• model_hash：模型的哈希值，可用于验证模型是否匹配策略。
• 拓扑：量化模型的拓扑
• node_conds : 哪些节点将被量化
• edge_conds : 哪些边将被量化
• bits : 每条边上的位数。
• 阈值：每个节点输出的阈值。
• 结果：实验结果
• sim_acc : 模拟模型的精度

 

搜索速度

实现

在得到最佳量化策略后：拓扑、比特、阈值实现模拟图，到低精度量化图，相当直截了当的。只需要用低精度整数运算，替换每条边上的 SimQ 运算。

调试

调试量化模型哪里出了问题，因为通常只知道最终的准确性很差。实现了inspect_graph_statistic逐层量化前后统计差异的功能，可以快速定位到哪里出错了。开发过程中，证明非常有帮助。

 

接口演示

from tvm import hago

# ideally we will have predefined description for x86, arm, gpu and vta

hardware = hago.create_sample_hardware()

strategy, sim_acc = hago.search_quantize_strategy(graph, hardware, dataset)

quantizer = hago.create_quantizer(graph, hardware, strategy)

simulated_graph = quantizer.simulate()

quantized_graph = quantizer.quantize()

当前状态

在 resnet18_v1 上获得了 68.7% 的初步结果，没有跳过第一个卷积层，只使用 2 的幂范围，不是 kl 距离，应该还有更多的改进空间。

参考链接：

https://discuss.tvm.apache.org/t/rfc-search-based-automated-quantization/5483

https://discuss.tvm.apache.org/t/quantization-story/3920