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
