本文是在学习TinyML L06时,因为视频没有新版slides中的内容,所以找来DoReFa-Net原文看看,看了很久才大致看懂,看了一些大佬的博客,但是他们没有描述更多细节,所以小白自己记录一下
内容总结写在前面
- 将BNN中的1-bit bit conv kernel扩展到多bit,计算复杂度为 O ( b i t w i d t h 1 × b i t w i d t h 2 ) O(bitwidth1 \times bitwidth2) O(bitwidth1×bitwidth2)
- 定义 q u a n t i z e k {quantize_{k}} quantizek 将浮点数输入 r i ∈ [ 0 , 1 ] r_{i} \in [0,1] ri∈[0,1] 量化到 k-bit输出 r o ∈ [ 0 , 1 ] r_{o}\in [0,1] ro∈[0,1]
- 定义权重量化、激活量化、梯度量化
- 实验表明量化所需要的位宽:梯度 > 激活 > 权重
- 权重和激活可以使用确定性量化,梯度需要随机量化
- 第一层和最后一层不量化
1 DoReFa-Net
1.1 比特卷积核
BNN中对于1-bit的点乘定义
x
⋅
y
=
b
i
t
c
o
u
n
t
(
a
n
d
(
x
,
y
)
)
,
x
i
,
y
i
∈
0
,
1
∀
i
{x}\cdot {y} = bitcount(and({x},{y})), x_{i},y_{i} \in {0,1} \forall i
x⋅y=bitcount(and(x,y)),xi,yi∈0,1∀i
假设
x
{x}
x 是 M-bit 定点数序列,
x
=
∑
m
=
0
M
−
1
c
m
(
x
)
2
m
{x}=\sum\limits_{m=0}^{M-1} c_{m}({x})2^{m}
x=m=0∑M−1cm(x)2m ,
y
{y}
y 是 K-bit 定点数序列,
y
=
∑
k
=
0
K
−
1
c
k
(
y
)
2
k
{y}=\sum\limits_{k=0}^{K-1} c_{k}({y})2^{k}
y=k=0∑K−1ck(y)2k 。其中,
(
c
m
(
x
)
)
m
=
0
M
−
1
(c_{m}({x}))_{m=0}^{M-1}
(cm(x))m=0M−1 和
(
c
k
(
y
)
)
k
=
0
K
−
1
(c_{k}({y}))_{k=0}^{K-1}
(ck(y))k=0K−1为bit vectors,点乘可以表示为:
x
⋅
y
=
∑
m
=
0
M
−
1
∑
k
=
0
K
−
1
2
m
+
k
b
i
t
c
o
u
n
t
[
a
n
d
(
c
m
(
x
)
,
c
k
(
y
)
)
]
{x}\cdot {y}=\sum_{m=0}^{M-1}\sum_{k=0}^{K-1}2^{m+k}bitcount[and(c_{m}({x}),c_{k}({y}))]
x⋅y=m=0∑M−1k=0∑K−12m+kbitcount[and(cm(x),ck(y))]
c
m
(
x
)
i
,
c
k
(
y
)
i
∈
{
0
,
1
}
,
∀
i
,
m
,
k
c_{m}({x})_{i},c_{k}({y})_{i}\in \{0,1\}, \forall i,m,k
cm(x)i,ck(y)i∈{0,1},∀i,m,k
此时计算复杂度为
M
∗
K
M*K
M∗K,即计算复杂度和位宽成正比
根据注释的符号解释
个人理解: x = { 10010 , 11011... } {x}=\{10010, 11011...\} x={10010,11011...} (在M=5)时,即 c m ( x ) i c_{m}({x})_{i} cm(x)i 下角标中的每个 i 代表一个M-bit定点数, c m ( ⋅ ) c_{m}(\cdot) cm(⋅) 表示一个定点数的第m位, x = ∑ m = 0 M − 1 c m ( x ) 2 m {x}=\sum\limits_{m=0}^{M-1} c_{m}({x})2^{m} x=m=0∑M−1cm(x)2m 表示:例如{10010}时,这个定点数=1x16+1x2=18
1.2 直通估计器STE
使用STE来规避梯度为0的问题
一个简单的例子是在Bernoulli分布采样中定义STE
其中,c为目标函数,由于Bernoulli采样是一个不可微分的过程,因此
∂
q
∂
p
\frac{\partial q}{\partial p}
∂p∂q 没有定义,因此反向传播不能由链式法则直接求得
∂
c
∂
p
\frac{\partial c}{\partial p}
∂p∂c
然而由于q和p具有相同的期望,可以使用 ∂ c ∂ p \frac{\partial c}{\partial p} ∂p∂c 对 ∂ c ∂ q \frac{\partial c}{\partial q} ∂q∂c 做近似,换句话说STE给出了一个自定义的 ∂ q ∂ p \frac{\partial q}{\partial p} ∂p∂q,即 ∂ q ∂ p = 1 \frac{\partial q}{\partial p}=1 ∂p∂q=1
本文定义
q
u
a
n
t
i
z
e
k
{quantize_{k}}
quantizek 将浮点数输入
r
i
∈
[
0
,
1
]
r_{i} \in [0,1]
ri∈[0,1] 量化到 k-bit输出
r
o
∈
[
0
,
1
]
r_{o}\in [0,1]
ro∈[0,1] ,其STE如下:
可以看到STE的输出是一个k-bit实数,而
r
o
r_{o}
ro是一个k-bit定点数,两个k-bit实数可以使用第一小节描述的定点数点乘(使用正确的缩放)
把这一段原文放在这里,我的理解是,
q
u
a
n
t
i
z
e
k
{quantize_{k}}
quantizek STE的输出是实数,而
q
u
a
n
t
i
z
e
k
{quantize_{k}}
quantizek的输出
r
o
r_{o}
ro为定点数,可以把STE的输出先按照定义的
q
u
a
n
t
i
z
e
k
{quantize_{k}}
quantizek转换成定点数,然后再使用比特卷积dot product
1.3 权重的低比特位宽量化
在之前的工作中,STE被用来二值化权重,例如在BNN中
在XNOR-Net中,权重在二值化后进行了缩放
在XNOR-Net中,缩放因子
E
F
(
∣
r
i
∣
)
{E}_{F}(|r_{i}|)
EF(∣ri∣)是对应每个权重的输出通道的绝对值取均值,引入这个scaling factor的可以增加权重值的表达范围,同时仍然能在前向传播卷积时利用比特卷积。
然而在反向传播计算梯度和权重之间的卷积时,channel-wise缩放因子的存在使得无法使用比特卷积核,因此,使用一个常量缩放因子来代替通道级缩放。
当k>1时,使用k-bit表示的权重,将STE应用在权重上:
量化到k-bit之前,使用
t
a
n
h
tanh
tanh来限制权重值得范围在-1~1之间,通过将量化的数值限制在[0, 1]之间,最大值是对整个层的权重而言的。然后通过
q
u
a
n
t
i
z
e
k
{quantize_{k}}
quantizek将其量化到[0, 1]之间的k-bit定点数,最后通过affine transform将结果映射到[-1, 1]之间
考虑对称性
需要注意的是,当k=1时,等式9不同于等式7,它提供了一个不同的二值化权重的方法,然而,论文发现在实验中这种区别不重要。
1.4 激活的低比特位宽量化
在BNN和XNOR-Net中,激活和权重采用相同的方法进行二值化。然而使用XNOR-Net中的方法无法得到很好的结果,且BNN中的二值化方法会在AlexNet这样的模型上损失很多精度。
因此,本文应用STE在输入激活r(each weight layer)。假设前一层的输出传递一个有界的激活函数h,确保
r
∈
[
0
,
1
]
r\in[0,1]
r∈[0,1],因此在DoReFa-Net中激活的量化简单定义为:
1.5 梯度的低比特位宽量化
在前面几节中描述了确定性量化可以产生低比特位宽的权重和激活,然而发现低位宽梯度必须采用随机量化。
要量化梯度到低比特位宽,需要注意梯度是无界的,并且可能具有比激活值更大的值范围。
其中,
d
r
=
∂
c
∂
r
dr=\frac{\partial c}{\partial r}
dr=∂r∂c 为一些层的输出 r 在反向传播时的梯度(对损失函数c的导数),最大致是对梯度张量的所有维度(除了batch size)的统计,然后在梯度上进行缩放将结果映射到0~1之间,并且在量化之后又缩放回去
为了进一步补偿潜在的由梯度量化所引入的bias,使用一个额外的噪声方程
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 服从 Uniform (-0.5,0.5)
σ服从Uniform(−0.5,0.5) 均匀分布,因此噪声可能具有和量化误差相同的幅值。实验发现人工噪声对达到好的performance至关重要。最终得到量化梯度的表达式
梯度的量化只在反向传播时进行,因此应用STE
2 DoReFa-Net训练算法
(哎呀少标记了一句,在12行的地方,量化梯度为Eq. 12)
以及这里的4 应该就算是后面PCAT拿来改进的地方
3 Source Code
注意这里返回的是three quantization functions
def get_dorefa(bitW, bitA, bitG):
"""
Return the three quantization functions fw, fa, fg, for weights, activations and gradients respectively
"""
def quantize(x, k):
n = float(2 ** k - 1)
@tf.custom_gradient
def _quantize(x):
return tf.round(x * n) / n, lambda dy: dy
return _quantize(x)
def fw(x):
if bitW == 32:
return x
if bitW == 1: # BWN
E = tf.stop_gradient(tf.reduce_mean(tf.abs(x)))
@tf.custom_gradient
def _sign(x):
return tf.where(tf.equal(x, 0), tf.ones_like(x), tf.sign(x / E)) * E, lambda dy: dy
return _sign(x)
x = tf.tanh(x)
x = x / tf.reduce_max(tf.abs(x)) * 0.5 + 0.5
return 2 * quantize(x, bitW) - 1
def fa(x):
if bitA == 32:
return x
return quantize(x, bitA)
def fg(x):
if bitG == 32:
return x
@tf.custom_gradient
def _identity(input):
def grad_fg(x):
rank = x.get_shape().ndims
assert rank is not None
maxx = tf.reduce_max(tf.abs(x), list(range(1, rank)), keep_dims=True)
x = x / maxx
n = float(2**bitG - 1)
x = x * 0.5 + 0.5 + tf.random_uniform(
tf.shape(x), minval=-0.5 / n, maxval=0.5 / n)
x = tf.clip_by_value(x, 0.0, 1.0)
x = quantize(x, bitG) - 0.5
return x * maxx * 2
return input, grad_fg
return _identity(x)
return fw, fa, fg
在Model类中调用fw,fa 实现了量化权重和激活
def new_get_variable(v):
name = v.op.name
# don't binarize first and last layer
if not name.endswith('W') or 'conv0' in name or 'fct' in name:
return v
else:
logger.info("Quantizing weight {}".format(v.op.name))
return fw(v)
def nonlin(x):
if BITA == 32:
return tf.nn.relu(x) # still use relu for 32bit cases
return tf.clip_by_value(x, 0.0, 1.0) # h()
def activate(x):
return fa(nonlin(x))
构建模型的时候应用了梯度的量化fg
4 参考内容
[低比特量化之DoreFa-Net理论与实践](低比特量化之DoreFa-Net理论与实践 - 知乎 (zhihu.com))