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

總結

以上是生活随笔為你收集整理的TVM量化代码解析的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。