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量化需要做的就是

l  标号1，图优化部分，具体做哪些图优化，就可自己选了，如算子融合，常量折叠。

l  标号2，整个量化的步骤，包括定义quant_passes，如果发现config设置，不需要伪量化，就是inference阶段了，就把realize加进去，否则，只需要annotate及calibrate，优化量化参数。

l  标号3，开始做量化了，将一个fp32的inference graph，转成int类型的inference graph，可参照第一张图。

l  标号4，把realize的graph，或者说对于一个op的前向推理的步骤，分成前中后三部分：

比如，conv2d，input_quantization -> input_quantization*weight_quantization(core function) -> ouput_dequantization，

每一个算子计算完后，都要dequant回去，很有可能某些算子没量化，还得用fp32。

最优解肯定是全部都量化掉，直接int32跑到底，TVM搞了个参数ensure_fully_integral，保证所有的算子都量化了。

 

 

参考链接：

https://blog.csdn.net/Artyze/article/details/108776522

https://www.freesion.com/article/3155559638/

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