论文(暂无投稿):https://arxiv.org/abs/1606.06160
源码(TensorFlow):https://github.com/tensorpack/tensorpack/tree/master/examples/DoReFa-Net
DoReFa-Net 概要
- 1 Motivation
- 2 Method
- 2.1 Using Bit Convolution Kernels In Low Bitwidth Neural Network
- 2.2 Straight-Through Estimator
- 2.3 Low Bitwidth Quantization of Weights
- 2.4 Low Bitwidth Quantization of Activations
- 2.5 Low Bitwidth Quantization of Gradients
- 2.6 The Algorithm for DoReFa-Net
- 2.7 First and the Last Layer
- 2.8 Reducing Run-time Memory Footprint by Fusing Nonlinear Function and Rounding
- 3 Experiment
- 个人思考
1 Motivation
- 先进的DCNN通常参数多、计算复杂度高,这些缺点不利于嵌入式设备中的应用。
- 以前的工作还没有成功地在向后传递期间将梯度量化为位宽小于8的离散值并保证性能不严重减少。
作者强调,现在的量化工作能很好地处理和加速卷积操作(全连接层)的前向传播过程,如公式 (1) 的点积型位运算(Bitwise Operation):
其中
x
\bm{\mathtt{x}}
x 和
y
\bm{\mathtt{y}}
y 为编码进制的比特向量,
b
i
t
c
o
u
n
t
bitcount
bitcount 用于计算一个比特向量里1的个数。当然当比特向量的每一位能表示
−
1
,
1
{-1, 1}
−1,1 时,位运算也可以表示为:
但是在反向传播过程中,梯度的格式仍为32-bit,于是训练过程中的梯度计算会耗占时间。
2 Method
作者注意到,虽然权重和激活可以确定量化,但梯度只能随机量化。
2.1 Using Bit Convolution Kernels In Low Bitwidth Neural Network
设 x \bm{\mathtt{x}} x 为 M M M-bit 定点整形的序列, 即 x = ∑ m = 0 M − 1 c m ( x ) 2 m \bm{\mathtt{x}}=\sum_{m=0}^{M-1}c_{m}(\bm{\mathtt{x}})2^m x=∑m=0M−1cm(x)2m 且 { c m ( x ) } m = 0 M − 1 \{c_{m}(\bm{\mathtt{x}})\}_{m=0}^{M-1} {cm(x)}m=0M−1 构成比特向量; y \bm{\mathtt{y}} y 为 K K K-bit 定点整形的序列, y = ∑ k = 0 K − 1 c k ( y ) 2 k \bm{\mathtt{y}}=\sum_{k=0}^{K-1}c_{k}(\bm{\mathtt{y}})2^k y=∑k=0K−1ck(y)2k 且 { c k ( y ) } k = 0 K − 1 \{c_{k}(\bm{\mathtt{y}})\}_{k=0}^{K-1} {ck(y)}k=0K−1 构成比特向量,那么对于低比特的定点整形来说有位操作:
上式的计算复杂度为
O
(
M
K
)
O(MK)
O(MK) ,即取决于位宽。
2.2 Straight-Through Estimator
为了解决量化过程几乎处处不可导,作者也引进了 STE 的思想。一个简单的例子是将 STE 作用于伯努利抽样 p ∈ [ 0 , 1 ] p \in [0, 1] p∈[0,1] :
就能使得原本不可导的前向传播,在反向传播时变得可导。其中
c
c
c 为目标函数。那么将STE应用于量化过程中,设实数输入
r
i
∈
[
0
,
1
]
r_{i}\in [0, 1]
ri∈[0,1] ,
k
k
k-bit 输出
r
o
∈
[
0
,
1
]
r_{o}\in [0, 1]
ro∈[0,1] ,就有
q
u
a
n
t
i
z
e
k
\bm{\mathsf{quantize}_{k}}
quantizek :
2.3 Low Bitwidth Quantization of Weights
当对 Weight 进行二值化操作时,作者定义:
其中
s
i
g
n
(
r
i
)
=
2
I
r
i
≥
0
−
1
sign(r_{i})=2 \mathbb{I}_{r_{i} \ge 0}-1
sign(ri)=2Iri≥0−1 能返回 -1 或 1,
E
F
(
∣
r
i
∣
)
\bm{E}_F(|r_{i}|)
EF(∣ri∣) 为整个 Weight 矩阵的绝对值的均值。而当对 Weight 进行
k
k
k-bit 表征量化(
k
>
1
k>1
k>1)时,有:
其中
t
a
n
h
tanh
tanh 将 Weight 钳制到范围
[
−
1
,
1
]
[-1, 1]
[−1,1] 。
t
a
n
h
(
r
i
)
2
m
a
x
(
∣
t
a
n
h
(
r
i
)
∣
)
+
1
2
\frac{tanh(r_{i})}{2max(|tanh(r_{i})|)}+\frac{1}{2}
2max(∣tanh(ri)∣)tanh(ri)+21 的范围为
[
0
,
1
]
[0, 1]
[0,1] ,需要注意此时的
m
a
x
max
max 是取整个 Weight 张量里的最大值。那么最后前向传播量化后的范围变成
[
−
1
,
1
]
[-1, 1]
[−1,1] 。
2.4 Low Bitwidth Quantization of Activations
设上一层的输出已经经过一个有界函数 h h h 的作用来保证 Activations r ∈ [ 0 , 1 ] r \in [0, 1] r∈[0,1] ,那么有量化:
2.5 Low Bitwidth Quantization of Gradients
作者发现,随机量化对梯度来说更有效。首先需要注意的是,梯度是无边界的,可能比激活具有更大的数值范围。于是作者初步设计方案如下:
其中
d
r
=
∂
c
∂
r
\mathrm{d} r=\frac{\partial c}{\partial r}
dr=∂r∂c,是某一层的输出
r
r
r 的后向传播梯度。而
m
a
x
max
max 取得是梯度张量的各个维度,出了 Mini-batch 维度(也就是说Mini-batch里的每个实例都有自己独立的缩放因子)。
为了补偿梯度量化所带来的潜在,作者又引入额外的噪声函数
N
(
k
)
=
σ
2
k
−
1
N(k)=\frac{\sigma}{2^k -1}
N(k)=2k−1σ ,
σ
∼
U
n
i
f
o
r
m
(
−
0.5
,
0.5
)
\sigma \sim Uniform(-0.5, 0.5)
σ∼Uniform(−0.5,0.5) 。这个人为的随机噪声反而起到关键作用,提高模型性能。于是最终梯度量化公式如下:
需要注意的是,上式只在反向传播中作用,作者定义了每个卷积层输出后接的 STE 构造方法:
2.6 The Algorithm for DoReFa-Net
2.7 First and the Last Layer
对于第一层,其输入通常为图像,其中可能包含 8-bit 特征。对于最后一层,由于输出的是 one-hot 估计向量,对 bit 也有一定要求。
在绝大多数情况下,作者的实验不对第一层及其输入、最后一层及其输出进行量化,但是这两层的输出会通过梯度量化器,对梯度进行量化。
2.8 Reducing Run-time Memory Footprint by Fusing Nonlinear Function and Rounding
实际上作者的网络并没有被完全量化,因为其中的一些中间过程涉及全精度的储存和计算,如 Algorithm 1 里的 Step 3、Step 4、Step 6都会储存全精度值。作者将这几步融合为同一操作
a
k
b
=
f
α
(
h
(
a
~
k
)
)
a_{k}^b=f_\alpha(h(\tilde{a}_k))
akb=fα(h(a~k)) ,就避免了中间过程的产生。类似的,Step 11、Step 12、Step 13 也能 Fuse 到一块。
在有 Pooling 层的情况下,上述操作会更复杂。且若 Pooling 是 Max-pooling 时,可以进行如下 Fuse 操作:
3 Experiment
Table 1 中,Model A 是作者自己构造的一个模型,输入为40x40 Image,有 7 个卷积层和 1 个全连接层。而 Model B,C,D 分别减少所有层的通道数的 50%,75%,87.5%。
Table 2 中,Initialize 意为使用预训练模型。
个人思考
需要注意,作者讨论模型结构时是忽略了 BN 层的。年限较早了,尝试对梯度进行量化是本文章最主要的创新点。
有一个问题,作者的 Weight 伪量化,量化前和量化后的范围是不一样的?这会对模型产生影响吗?