BN折叠及其量化

本文介绍量化过程中的BN折叠;

以下图片引用自 Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference

BN折叠

首先介绍BN折叠;BN折叠源自 BN模块在训练阶段和推断阶段的实现差异;

如下图所示:以conv+BN为例,图1为训练阶段的BN,图2为推断阶段的BN

BN折叠及其量化_第1张图片 图1:训练阶段的BN

BN折叠及其量化_第2张图片 图2:推断阶段的BN,BN被折叠刀CONV中

可以证明这样的折叠是合理的,注意图2中的w指的是卷积核;

\\ O_f = \gamma * \frac{\sum(W * I) + b_{conv} - \mu}{\sigma} + b_{bn} \\ O_f =\sum( \frac{\gamma}{\sigma}*W*I) + \frac{\gamma}{\sigma}*(b_{conv} - \mu) + b_{bn} \\ => W' = \frac{\gamma}{\sigma}*W \\ => b'_{conv} = \frac{\gamma}{\sigma}*(b_{conv} - \mu) + b_{bn} \\ => b'_{conv} = \frac{\gamma}{\sigma}*(- \mu) + b_{bn} = b_{bn}- \frac{\gamma\mu}{\sigma} \space\ \space\ \space\ if \space\ \space\ b_{conv} == 0


量化中的BN折叠

对量化来说,尤其是量化感知(伪量化),我们需要保持训练阶段的伪量化和上述图2的推断阶段的量化实现是一致的,因此对训练阶段的CONV+BN二阶段步骤进行折叠,如下图3:

BN折叠及其量化_第3张图片 图3:训练阶段的BN折叠改造,思想源自图2所示实现
​​​

图3的设计思想源自图2,基本思想是训练阶段求\mu/\sigma的EMA(即图中moment及其前后步骤)拿掉后要和图2一致;

根据图3所示BN折叠,就可以对训练阶段该折叠做后训练量化或者伪量化,如下图4

BN折叠及其量化_第4张图片 图4:BN折叠后训练阶段的(伪)量化过程


代码实现

以下代码块来自github distiller,没有完整的上下文实现,感兴趣的读者可以移步阅读

图5:BN折叠量化实现


总结

  • 以上简要介绍了BN折叠及其量化,思想比较巧妙,感兴趣的读者可以根据文中提到的文献进一步深入阅读,谢谢
  • 理解有误的地方欢迎指正

你可能感兴趣的:(技术问题,#,量化,#,深度学习,深度学习,人工智能)