Paper name
DOREFA-NET: TRAINING LOW BITWIDTH CONVOLUTIONAL NEURAL NETWORKS WITH LOW BITWIDTH GRADIENTS
Paper Reading Note
URL: https://arxiv.org/pdf/1606.06160.pdf
TL;DR
- 提供了低比特权重、特征、梯度的在线量化训练方案
- 量化位宽的重要性排序为:梯度>特征>权重,所以实验中的配比建议为 1w4f6g(1bit权重,4bit特征,6bit梯度),配比的关系对应于 DoReFa 的论文名
- 首个成功在低于 8bit 梯度下得到较高精度的在线量化训练方案
- SVHN 和 ImageNet 数据集上基于 AlexNet 的实验结果证明精度可以与浮点模型相比
Introduction
- 之前的量化方法对于梯度的量化位宽没有低于 8bit 的(最少是 10bit,BNN 和 XNOR-Net 都是全精度梯度训练,导致大部分训练时间消耗在反传环节),这里提出一种能在低于 8bit 位宽梯度下的训练方案
- 这种梯度低比特量化的方案后面可能能支持在 FPGA、ASIC 上的低功耗训练
- 量化位宽的重要性排序为:梯度>特征>权重,配比的关系对应于 DoReFa 的论文名(1-bit weights, 1-bit activations and 2-bit gradients can lead to 93% accuracy on SVHN dataset)
Dataset/Algorithm/Model/Experiment Detail
实现方式
1. STRAIGHT-THROUGH ESTIMATOR(STE,直通滤波器)
- STE是对不可微分算子在计算反向传播梯度时的近似方法,比如量化操作不可微(量化函数一般是阶梯函数,处处的梯度为0),STE就假装量化没有发生,将梯度直接前传
- 比如对于伯努利采样函数, c是目标函数,伯努利分布不可微,
∂
(
c
)
∂
(
p
)
\frac{\partial(c)}{\partial(p)}
∂(p)∂(c) 没有很好的定义,因此不能使用链式规则直接从正向过程构造反向过程,因为
p
p
p 和
q
q
q 是期望相等的,所以使用能够被正常定义的
∂
(
c
)
∂
(
q
)
\frac{\partial(c)}{\partial(q)}
∂(q)∂(c) 来近似
∂
(
c
)
∂
(
p
)
\frac{\partial(c)}{\partial(p)}
∂(p)∂(c)
- 又比如对于如下量化函数,将 [0, 1] 之间的全精度数量化到 [0,1] 之间的 k-bit 量化数,注意前传的量化函数不唯一,可以是不同的量化形式
- STE 的问题:
- loss curve会波动
- 由于量化函数实际上和y=x这条直线并不一致,而我们用后者的梯度来代替前者,这就会导致梯度的不匹配
2. 权重量化
-
对于 1bit 权重情况下量化方法如下,其中的scaling factor: E ( ∣ r i ∣ ) E(|r_{i}|) E(∣ri∣) 代表该层的权重绝对值的整体期望 ,这个 scaling factor 的作用是增加权重的取值范围;与 XNOR-Net 量化方法的区别是 XNOR-Net 中的 E ( ∣ r i ∣ ) E(|r_{i}|) E(∣ri∣) 代表的是该层权重的每个 channel 的绝对值的期望,但是这里使用逐 channel 的 scaling factor 会导致计算在反向传播过程中,当计算梯度和权重之间的卷积时无法利用位卷积核
-
对于k-bit(k>1)的权重量化
- 其中
q
u
a
n
t
i
z
e
k
quantize_{k}
quantizek 为量化函数,将 [0, 1] 之间的全精度浮点数聚类到 [0, 1] 区间的定点数,即将 [0, 1] 的浮点数聚类到 [0, 1] 中的
2
k
2^{k}
2k 个等距分割点上
- 其中
q
u
a
n
t
i
z
e
k
quantize_{k}
quantizek 为量化函数,将 [0, 1] 之间的全精度浮点数聚类到 [0, 1] 区间的定点数,即将 [0, 1] 的浮点数聚类到 [0, 1] 中的
2
k
2^{k}
2k 个等距分割点上
-
权重量化源码
def uniform_quantize(k): class qfn(torch.autograd.Function): @staticmethod def forward(ctx, input): if k == 32: out = input elif k == 1: out = torch.sign(input) else: n = float(2 ** k - 1) out = torch.round(input * n) / n return out @staticmethod def backward(ctx, grad_output): # STE (do nothing in backward) grad_input = grad_output.clone() return grad_input return qfn().apply class weight_quantize_fn(nn.Module): def __init__(self, w_bit): super(weight_quantize_fn, self).__init__() assert w_bit <= 8 or w_bit == 32 self.w_bit = w_bit self.uniform_q = uniform_quantize(k=w_bit) def forward(self, x): if self.w_bit == 32: weight_q = x elif self.w_bit == 1: E = torch.mean(torch.abs(x)).detach() weight_q = self.uniform_q(x / E) * E else: weight = torch.tanh(x) max_w = torch.max(torch.abs(weight)).detach() weight = weight / 2 / max_w + 0.5 weight_q = max_w * (2 * self.uniform_q(weight) - 1) return weight_q
3. 特征量化
- 特征是简单的直接量化处理
4. 梯度量化
-
梯度量化方法
梯度对于量化更为敏感,所以前人的低比特梯度量化一般不 work,为了进一步补偿梯度量化引入的潜在偏差,这里加上随机均匀噪声从而增加鲁棒性,噪声的量级和量化误差类似:
-
一般板端部署模型时不需要在板端进行训练,所以一般都是使用全精度的梯度在服务器端进行模型训练,然后在板端部署,避免梯度量化产生的掉点
5. 整体算法流程
- 需要注意的是这些低比特的层之间会有高比特精度的 affine 层连接,从而节省了 conv 层等高功耗层的计算量
6. 不同层的量化位宽处理
- 第一层因为量化掉点明显并且第一层的 channel 数一般较少,所以一般第一层权重不量化,仅做特征量化
- 最后一层考虑到输出 channel 或者 fc 节点数较少,不量化的影响较小,为了保证精度一般最后一层的权重和特征都不做量化
实验结果
- 可以看出在梯度的位宽有保证的情况下相比对浮点模型的掉点较少
Thoughts
- 低比特量化的经典作品,梯度量化对于板端训练的意义重大