TVM量化代码解析
2021/10/31 6:39:37
本文主要是介绍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
这篇关于TVM量化代码解析的文章就介绍到这儿,希望我们推荐的文章对大家有所帮助,也希望大家多多支持为之网!
- 2024-05-01为什么公共事业机构会偏爱 TiDB :TiDB 数据库在某省妇幼健康管理系统的应用
- 2024-04-26敏捷开发:想要快速交付就必须舍弃产品质量?
- 2024-04-26静态代码分析的这些好处,我竟然都不知道?
- 2024-04-26你在测试金字塔的哪一层?(下)
- 2024-04-26快刀斩乱麻,DevOps让代码评审也自动起来
- 2024-04-262024年最好用的10款ER图神器!
- 2024-04-2203-为啥大模型LLM还没能完全替代你?
- 2024-04-2101-大语言模型发展
- 2024-04-17基于SpringWeb MultipartFile文件上传、下载功能
- 2024-04-14个人开发者,Spring Boot 项目如何部署