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

上一篇:vue学习第一天


下一篇:[ARC124E] Pass to Next