文章目录
Training Quantized Neural Networks with a Full-precision Auxiliary Module
Introduction
作者认为,量化模型性能下降的主要原因在于在训练过程中,quantizer是不可导的,无法直接使用SGD来优化网络。目前的解决方案可分为两类,一类是直接使用STE,或者是使用一个“松弛”的quantizer,另一类则是用全精度网络来指导量化网络训练,如knowledge distillation。作者借鉴第二类方法,提出了一种全精度网络+量化网络的训练方法,在前向推理阶段,去掉全精度网络,可得到一个精度较高的量化网络。
Method
作者提出的网络训练框架如上图所示。蓝色图部分表示为
F
F
F,为量化网络。粉红色框记作
H
H
H,表示全精度网络,即为auxiliary module。在训练阶段,两块的loss都会计算,全精度网络和量化网络同时进行训练,在测试阶段,
H
H
H将会被舍弃,只保留
F
F
F。
在上图中,
H
H
H是由一系列的adaptor和aggregators组成,这个辅助性的
H
H
H从
F
F
F中接收
P
P
P个特征图的输出
{
O
p
}
p
=
1
P
\left\{\mathbf{O}_{p}\right\}_{p=1}^{P}
{Op}p=1P,,假定
{
B
1
,
…
,
B
P
}
\left\{B_{1}, \ldots, B_{P}\right\}
{B1,…,BP}表示生成feature map的block,对于
H
H
H中的第
p
p
p个输入,作者使用一个可训练的adaptor
ϕ
p
(
⋅
)
\phi_{p}(\cdot)
ϕp(⋅) 对得到的特征图
O
p
\mathbf{O}_{p}
Op进行加工,并生成新的特征图
ϕ
p
(
O
p
)
\phi_{p}\left(\mathbf{O}_{p}\right)
ϕp(Op)。作者这么做的动机为compensate the distribution discrepancy between the low-precision model and full-precision model。It ensures the quantized activations
{
O
p
}
p
=
1
P
\left\{\mathbf{O}_{p}\right\}_{p=1}^{P}
{Op}p=1P to be compatible to the full-precision calculation in
H
H
H。在具体实现的时候,作者使用一个
1
×
1
1\times 1
1×1卷积和一个BN层进行实现。
在
H
H
H中,从adaptor输入的特征图
ϕ
p
(
O
p
)
\phi_{p}\left(\mathbf{O}_{p}\right)
ϕp(Op)和
H
H
H自身的特征图
g
p
\boldsymbol{g}_{p}
gp相加,然后送入
ReLU
(
⋅
)
\operatorname{ReLU}(\cdot)
ReLU(⋅)当中,得到
g
p
=
ReLU
(
ϕ
p
(
O
p
)
+
g
p
−
1
)
\boldsymbol{g}_{p}=\operatorname{ReLU}\left(\phi_{p}\left(\mathbf{O}_{p}\right)+\boldsymbol{g}_{p-1}\right)
gp=ReLU(ϕp(Op)+gp−1)关于该网络的优化,对于训练样本
{
x
i
,
y
i
}
i
=
1
N
\left\{\boldsymbol{x}_{i}, \boldsymbol{y}_{i}\right\}_{i=1}^{N}
{xi,yi}i=1N,主干网络
F
F
F和混合精度网络
F
∘
H
F \circ H
F∘H,训练的损失函数为
min
{
θ
F
,
θ
H
}
∑
i
=
1
N
L
(
F
(
x
i
;
θ
F
)
,
y
i
)
+
L
a
u
x
(
(
F
∘
H
)
(
x
i
;
θ
H
,
θ
F
)
,
y
i
)
\begin{aligned} \min _{\left\{\boldsymbol{\theta}^{F}, \boldsymbol{\theta}^{H}\right\}} \sum_{i=1}^{N} \mathcal{L}\left(F\left(\boldsymbol{x}_{i} ; \boldsymbol{\theta}^{F}\right), \boldsymbol{y}_{i}\right) +\mathcal{L}_{a u x}\left((F \circ H)\left(\boldsymbol{x}_{i} ; \boldsymbol{\theta}^{H}, \boldsymbol{\theta}^{F}\right), \boldsymbol{y}_{i}\right) \end{aligned}
{θF,θH}mini=1∑NL(F(xi;θF),yi)+Laux((F∘H)(xi;θH,θF),yi)式中,
θ
F
\boldsymbol{\theta}^{F}
θF和
θ
H
\boldsymbol{\theta}^{H}
θH分别表示
F
F
F和
H
H
H网络参数,
L
\mathcal{L}
L表示目标损失函数,
L
a
u
x
\mathcal{L}_{a u x}
Laux表示辅助网络的loss,在分类任务中,这两个loss都设置为交叉熵损失函数。从上式可以看出,
θ
F
\boldsymbol{\theta}^{F}
θF是由
F
F
F和
F
∘
H
F \circ H
F∘H所共享,在求导的时候,作者说将两路的梯度取一个平均(不太理解为什么不直接求导?),作者给的理由是梯度方向更为准确,论文后面的一句话也没看懂,In other words, the full-precision module
H
H
H provides direct gradient for
F
F
F using weight sharing during back-propagation。整个训练算法的流程如下如下所示,如下图来看,作者确实是对梯度取了平均,即对loss取了平均。
作者还讨论了这种方法和其他方法的优劣,如和knowledge distillation或者是用中间特征图来计算分类loss的区别,具体可参见原论文。
Experiment
在分类实验的时候,作者将第一个卷积层和最后的FC层量化到8-bit,其他层量化到超低bit。结果如下图
从这个结果来看,如果单看绝对精度,2-bit量化没有特别高,而且还用DOReFa量化,效果不是很好,但是和自身相比的话,确实也说明了辅助网络的存在能够提升量化性能,同时,作者还在检测任务上进行了实验,4-bit量化基本不掉点,总来的说,应该还是有一些效果的,对于在其他量化方法的性能提升,就需要具体的实验验证,不过作者提供的这种训练的思路还是可取的,缺点就是会让训练的时候的网络变大,训练变慢。