TVM系列 - 量化

TVM量化原理

关于量化的方式其实已经有足够的文章去了解目前最主流的两种：离线量化及训练时量化（大家应该能理解，其实就是伪量化），而tvm的作者之一也对这些文章有一个初步的汇总：tvm作者关于量化方式的讨论，这部分不做过多讨论，直接进入tvm量化的部分

TVM量化现状

事实上，可能tvm团队将主要精力放到了autotvm，以框架整体的优化上（毕竟这部分才是tvm的特点呀），tvm的单纯量化效果并不理想，不仅精度有所下降，正常的，但是速度还有可能会变慢。这点还有待我自己亲自去验证：
TVM8bit测试
不过再配合上优秀的AutoTVM，这些都不是问题了。
当然，随着TVM的不断发展，相信TVM也会逐渐把量化这块，不仅是精度速度的提升，更是针对更多bit情况的支持。

TVM量化原理介绍

TVM的量化总体可以分为三个阶段：
Annotation：annotation会将TVM原本的pass进行重写，根据不同算子（var）类型插入伪量化
Calibration：熟悉量化的同学应该都知道，weights的量化所需要的的(f_min, f_max)并不会变，就直接求当前张量中的最大最小值。但是activation的量化(f_min, f_max)是需要随着训练样本的变化不断维护一个整体的值得，就像bn层scale和bias一样，一般的方式是通过exponential moving averages（EMA），而calibration的作用就是调整这个(f_min, f_max)
Realization:这个就非常好理解了，根据伪量化训练得到的参数，进行实质的量化过程

直观一点的看TVM整个量化流程的话，整体还是延续了tensorflow伪量化的路线，为pass中每一个中间var插入一个伪量化节点，而具体伪量化的方法与tensorflow稍有差异，tensorflow跑的是非对称量化，而TVM目前默认跑的是对称量化。
对比看下tensorflow量化与TVM量化的公式：

Tensorflow
TVM

让我们给他转换一下：

这样看上去是不是就舒服多了，其实他们做的都是同样的事情，将从fp32到int8的误差，包括rounding的误差及clamp的误差，在训练过程中即考虑进去，只不过确实对称量化会看上去有点膈应。

上图是TVM为conv2d算子插入伪量化的示意图，需要注意的是，

1. TVM插入伪量化的时候是要对input，weights,以及activation进行区分的，这也对应了我前面所说的，这两个的参数优化是不同的。

2. 由于伪量化无法拿到实际跑得模型的数据分布和数据类型的，具体要等到真正跑的时候才知道，所以到时候会不会有溢出的风险伪量化是不知道的，这是非常危险的。事实上，真正成熟的量化算法，都是要对量化过程中每个量化出来的张量进行一个缩放，保证不会有溢出的风险。

3. 上图就是TVM量化的实际操作过程，首先根据模型和要部署的硬件，选择要量化的ops，也就是Topology，以及量化的目标bit数，这里要量化到几bit并不是人为确定（当然肯定也可以人为写死），而是一个不断搜索的过程；然后根据一个小的calibration的图片集获得threshold，将threshold-bit-topology确定的量化模型去测试下，然后进行下一个标定循环 。

4. 下面来聊聊伪量化在训练标定过程中都要解决哪些问题。一般认为量化的损失来自于两个方面：

   • 取整损失，例如r = [6.8, 7.2, -0.6], scale = (7.2+0.6)/127 = 0.061417, q1 = 7.2/scale = 117.23,那么他的量化值就是117，有了0.23的损失
   • 截断损失 ，因为scale是取最优区间，那么边界的点势必会有超过最大量化值得情况，那么就要截断到最大量化值，就有了截断损失

   那么就要针对这两种损失，将这两种损失都加入到伪量化中去，在训练阶段就将误差消化掉。我们看下TVM的伪量化代码如何实现：

@_op.register_compute("relay.op.annotation.simulated_quantize")
def simulated_quantize_compute(attrs, inputs, out_type):
    """Compiler for simulated_quantize."""
    assert len(inputs) == 4
    assert attrs.sign
    assert attrs.rounding == "round"

    data, scale, clip_min, clip_max = inputs

    if attrs.kind == QAnnotateKind.IDENTITY:
        return [topi.identity(data)]
        
    scaled_data = topi.divide(data, scale)
    # simulate saturated error
    clipped_data = topi.maximum(topi.minimum(scaled_data, clip_max), clip_min)
    # simulate rounding error
    round_data = topi.round(clipped_data)

    # recover data
    rdata = topi.multiply(round_data, scale)
    return [rdata]

TVM量化代码解析

TVM的量化非常方便，即插即用，前面说过，TVM量化的话就是使用加入了伪量化后的pass替代原来的pass，我们看一个官方提供的例子：

def test_mul_rewrite():
    """a test case where rhs of mul is not constant"""
    data = relay.var("data", shape=(1, 16, 64, 64))
    multiplier = relay.sigmoid(relay.var("data", shape=(1, 16, 1, 1)))
    conv = relay.nn.conv2d(data, relay.var("weight"),
                           kernel_size=(3, 3),
                           padding=(1, 1),
                           channels=16)
    act = relay.nn.relu(data=conv)

    quantize_and_build(act * multiplier)

    pool = relay.nn.global_avg_pool2d(data=act)

    quantize_and_build(act * pool)

入口就是函数：

def quantize_and_build(out):
    f = relay.Function(relay.analysis.free_vars(out), out)
    mod, params = testing.create_workload(f)

    with relay.quantize.qconfig(skip_conv_layers=[]):
        qmod = relay.quantize.quantize(mod, params)
        
    relay.build(qmod, "llvm", params=params)
    return qmod

可以看到这里，调用的就是relay.quantize.quantize函数，因为这个函数实在太长了，我就只是放上来主体部分，

	

 1. mod = prerequisite_optimize(mod, params)

 2.   calibrate_pass = tvm.transform.module_pass(
        calibrate(dataset), opt_level=1,
        name="QuantizeCalibrate")
    quant_passes = [partition(),
                    annotate(),
                    calibrate_pass]
    if not current_qconfig().do_simulation:
        quant_passes.append(realize())
    quant_passes.append(_transform.FoldConstant())
    quantize_seq = tvm.transform.Sequential(quant_passes)
    with tvm.transform.PassContext(opt_level=3,
                                   required_pass=["QuantizeAnnotate",
                                                  "QuantizeCalibrate",
                                                  "QuantizeRealize"]):
 3.       with quantize_context():
            mod = quantize_seq(mod)

 4. q_cfg = current_qconfig()
    assert q_cfg.partition_conversions in ['disabled', 'enabled', 'fully_integral']
    if q_cfg.partition_conversions != 'disabled':
        quantized_dtypes = {q_cfg.dtype_input, q_cfg.dtype_weight, q_cfg.dtype_activation}
        ensure_fully_integral = q_cfg.partition_conversions == 'fully_integral'
        return partition_conversions(mod, quantized_dtypes, ensure_fully_integral)

从代码中可以看到，TVM量化首先需要做的就是
标号1，就是图优化部分，具体做哪些图优化就可以自己选了，比如算子融合，常量折叠这些。（顺便说下，关于TVM的下一篇文章应该就是图优化的部分了哦，先从基于支配树的算子融合开始）
接下来，就是
标号2那里定义，整个量化的步骤，包括定义quant_passes,如果发现config中设置不需要伪量化，也就是现在是inference阶段了，那就把realize加进去，否则只需要annotate及calibrate去优化量化参数就够了。
标号3就是说，现在开始做量化了，将一个fp32的inference graph转成int类型的inference graph，具体可以参照第一张图，没什么好说的。
标号4部分，意思就是，把realize的graph，或者说对于一个op的前向推断的步骤分成前中后三部分：比如conv2d，input_quantization -> input_quantization*weight_quantization(core function) -> ouput_dequantization，之所以要每一个算子计算完之后都要dequant回去，是因为很有可能某些算子没量化呀，你不还得用fp32嘛 😦
不过最优解肯定是全部都量化掉，直接int32跑到底，所以TVM搞了个参数ensure_fully_integral来保证所有的算子都被量化了

本文所有内容都是基于自己理解所写，肯定有不对的地方，欢迎大家交流讨论