TVM系列 - 量化

TVM量化原理

关于量化的方式其实已经有足够的文章去了解目前最主流的两种:离线量化及训练时量化(大家应该能理解,其实就是伪量化),而tvm的作者之一也对这些文章有一个初步的汇总:tvm作者关于量化方式的讨论,这部分不做过多讨论,直接进入tvm量化的部分

TVM量化现状

事实上,可能tvm团队将主要精力放到了autotvm,以框架整体的优化上(毕竟这部分才是tvm的特点呀),tvm的单纯量化效果并不理想,不仅精度有所下降,正常的,但是速度还有可能会变慢。这点还有待我自己亲自去验证:
TVM8bit测试
不过再配合上优秀的AutoTVM,这些都不是问题了。
当然,随着TVM的不断发展,相信TVM也会逐渐把量化这块,不仅是精度速度的提升,更是针对更多bit情况的支持。

TVM量化原理介绍

TVM的量化总体可以分为三个阶段:
Annotation:annotation会将TVM原本的pass进行重写,根据不同算子(var)类型插入伪量化
Calibration:熟悉量化的同学应该都知道,weights的量化所需要的的(f_min, f_max)并不会变,就直接求当前张量中的最大最小值。但是activation的量化(f_min, f_max)是需要随着训练样本的变化不断维护一个整体的值得,就像bn层scale和bias一样,一般的方式是通过exponential moving averages(EMA),而calibration的作用就是调整这个(f_min, f_max)
Realization:这个就非常好理解了,根据伪量化训练得到的参数,进行实质的量化过程
TVM量化流程

直观一点的看TVM整个量化流程的话,整体还是延续了tensorflow伪量化的路线,为pass中每一个中间var插入一个伪量化节点,而具体伪量化的方法与tensorflow稍有差异,tensorflow跑的是非对称量化,而TVM目前默认跑的是对称量化。
对比看下tensorflow量化与TVM量化的公式:

Tensorflow 将会带来全新的写作体验;
TVM
在这里插入图片描述
让我们给他转换一下:
在这里插入图片描述

这样看上去是不是就舒服多了,其实他们做的都是同样的事情,将从fp32到int8的误差,包括rounding的误差及clamp的误差,在训练过程中即考虑进去,只不过确实对称量化会看上去有点膈应。
TVM conv层插入伪量化示例
上图是TVM为conv2d算子插入伪量化的示意图,需要注意的是,

  1. TVM插入伪量化的时候是要对input,weights,以及activation进行区分的,这也对应了我前面所说的,这两个的参数优化是不同的。

  2. 由于伪量化无法拿到实际跑得模型的数据分布和数据类型的,具体要等到真正跑的时候才知道,所以到时候会不会有溢出的风险伪量化是不知道的,这是非常危险的。事实上,真正成熟的量化算法,都是要对量化过程中每个量化出来的张量进行一个缩放,保证不会有溢出的风险。tvm量化算法更新

  3. 上图就是TVM量化的实际操作过程,首先根据模型和要部署的硬件,选择要量化的ops,也就是Topology,以及量化的目标bit数,这里要量化到几bit并不是人为确定(当然肯定也可以人为写死),而是一个不断搜索的过程;然后根据一个小的calibration的图片集获得threshold,将threshold-bit-topology确定的量化模型去测试下,然后进行下一个标定循环 。

  4. 下面来聊聊伪量化在训练标定过程中都要解决哪些问题。一般认为量化的损失来自于两个方面:

    • 取整损失,例如r = [6.8, 7.2, -0.6], scale = (7.2+0.6)/127 = 0.061417, q1 = 7.2/scale = 117.23,那么他的量化值就是117,有了0.23的损失
    • 截断损失 ,因为scale是取最优区间,那么边界的点势必会有超过最大量化值得情况,那么就要截断到最大量化值,就有了截断损失

    那么就要针对这两种损失,将这两种损失都加入到伪量化中去,在训练阶段就将误差消化掉。我们看下TVM的伪量化代码如何实现:

@_op.register_compute("relay.op.annotation.simulated_quantize")
def simulated_quantize_compute(attrs, inputs, out_type):
    """Compiler for simulated_quantize."""
    assert len(inputs) == 4
    assert attrs.sign
    assert attrs.rounding == "round"

    data, scale, clip_min, clip_max = inputs

    if attrs.kind == QAnnotateKind.IDENTITY:
        return [topi.identity(data)]
        
    scaled_data = topi.divide(data, scale)
    # simulate saturated error
    clipped_data = topi.maximum(topi.minimum(scaled_data, clip_max), clip_min)
    # simulate rounding error
    round_data = topi.round(clipped_data)

    # recover data
    rdata = topi.multiply(round_data, scale)
    return [rdata]

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量化首先需要做的就是
标号1,就是图优化部分,具体做哪些图优化就可以自己选了,比如算子融合,常量折叠这些。(顺便说下,关于TVM的下一篇文章应该就是图优化的部分了哦,先从基于支配树的算子融合开始)
接下来,就是
标号2那里定义,整个量化的步骤,包括定义quant_passes,如果发现config中设置不需要伪量化,也就是现在是inference阶段了,那就把realize加进去,否则只需要annotate及calibrate去优化量化参数就够了。
标号3就是说,现在开始做量化了,将一个fp32的inference graph转成int类型的inference graph,具体可以参照第一张图,没什么好说的。
标号4部分,意思就是,把realize的graph,或者说对于一个op的前向推断的步骤分成前中后三部分:比如conv2d,input_quantization -> input_quantization*weight_quantization(core function) -> ouput_dequantization,之所以要每一个算子计算完之后都要dequant回去,是因为很有可能某些算子没量化呀,你不还得用fp32嘛 😦
不过最优解肯定是全部都量化掉,直接int32跑到底,所以TVM搞了个参数ensure_fully_integral来保证所有的算子都被量化了

本文所有内容都是基于自己理解所写,肯定有不对的地方,欢迎大家交流讨论

评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值