基于可训练Step-size的低比特量化——LSQ: Learned Step-size Quantization

Paper地址:https://arxiv.org/abs/1902.08153

GitHub地址 (PyTorch):GitHub - zhutmost/lsq-net: Unofficial implementation of LSQ-Net, a neural network quantization framework

基本量化设置

  • 计算结点伪量化:
  1. Weight跟Activation都采用Per-tensor量化;
  2. Scaling factor (Paper标记为Step size)作为量化参数,是可训练变量;
  3. 另外,针对TensorRT、MNN等推理引擎,Weight通常执行Per-channel量化,Activation执行Per-tensor量化;为了加快量化训练收敛,Activation的量化参数(可训练)可借助KL量化、或PyTorch observer量化予以初始化,Weight的量化参数则根据absmax方法在线更新;

  • 量化计算公式:

Step-size梯度求解

  • Loss关于step-size的梯度推导:

  1. 基于STE (Straight Through Estimate),直接Bypass掉Round结点的梯度反传(直接近似):

  • 与其他方法梯度函数的对比:
  1. QIL:https://arxiv.org/abs/1808.05779
  2. PACT:https://arxiv.org/abs/1805.06085
  3. 相比较而言,LSQ无论在不同Bin之间的Transition位置、还是量化范围的Clip位置,都有较明显的梯度数值,因而更有利Step-size的参数更新:

  • 梯度的scale:
    • 当量化比特数增加时,step-size会变小,以确保更为精细的量化;而当量化比特数减少时,step-size会变大。为了让step-size的参数更新,能够适应量化比特数的调整,需要将step-size的梯度乘以一个scale系数,例如针对weight:

  • 再奉上TensorFlow版本的LSQ子图构造代码:
def lsq_grad_scale(inputs, scale_var, max_value=127.0):
        """Gradient scale computation for LSQ."""
        tensor_shape = array_ops.shape(inputs)[1:]
        grad_scale = (
            math_ops.sqrt(
                math_ops.cast(math_ops.reduce_prod(tensor_shape), dtypes.float32)
                * max_value
            )
            / self._grad_scale
        )

        grad_scale = array_ops.stop_gradient(grad_scale)

        # stop gradient backward
        scale_div = math_ops.divide(scale_var, grad_scale)
        scale_diff = array_ops.stop_gradient(scale_var - scale_div)

        step_size = scale_diff + scale_div

        return step_size

量化训练

  • 伪量化训练时,从预训练模型加载FP32参数,学习率按余弦曲线衰减;通常第一层与最后一层固定为INT8量化,有些实际应用场景中亦可固定为FP32(量化敏感层)。

实验结果

  • 2
    点赞
  • 4
    收藏
  • 打赏
    打赏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
©️2022 CSDN 皮肤主题:书香水墨 设计师:CSDN官方博客 返回首页
评论

打赏作者

Law-Yao

你的鼓励将是我创作的最大动力

¥2 ¥4 ¥6 ¥10 ¥20
输入1-500的整数
余额支付 (余额:-- )
扫码支付
扫码支付:¥2
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值