MIT-TinyML学习笔记【6】DoReFa-Net

本文是在学习TinyML L06时，因为视频没有新版slides中的内容，所以找来DoReFa-Net原文看看，看了很久才大致看懂，看了一些大佬的博客，但是他们没有描述更多细节，所以小白自己记录一下

内容总结写在前面

• 将BNN中的1-bit bit conv kernel扩展到多bit，计算复杂度为$O(bitwidth1\times bitwidth2)$
• 定义$quantiz{e}_k$ 将浮点数输入 $r_i\in [0,1]$ 量化到 k-bit输出 $r_o\in [0,1]$
• 定义权重量化、激活量化、梯度量化
• 实验表明量化所需要的位宽：梯度 > 激活 > 权重
• 权重和激活可以使用确定性量化，梯度需要随机量化
• 第一层和最后一层不量化

1 DoReFa-Net

1.1 比特卷积核

BNN中对于1-bit的点乘定义
$$x\cdot y=bitcount(and(x,y)),x_i,y_i\in 0,1\forall i$$
假设 $x$ 是 M-bit 定点数序列，$x=\sum _{m=0}^{M-1}{c}_m(x)2^m$ ，$y$ 是 K-bit 定点数序列，$y=\sum _{k=0}^{K-1}{c}_k(y)2^k$ 。其中，$(c_m(x){)}_{m=0}^{M-1}$ 和$(c_k(y){)}_{k=0}^{K-1}$为bit vectors，点乘可以表示为：

$$x\cdot y=\sum _{m=0}^{M-1}\sum _{k=0}^{K-1}{2}^{m+k}bitcount[and(c_m(x),c_k(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\ast K$，即计算复杂度和位宽成正比

根据注释的符号解释

个人理解： x = { 10010 , 11011... } {x}=\{10010, 11011...\} x={10010,11011...} (在M=5)时，即 $c_m(x{)}_i$ 下角标中的每个 i 代表一个M-bit定点数，$c_m(\cdot )$ 表示一个定点数的第m位，$x=\sum _{m=0}^{M-1}{c}_m(x)2^m$ 表示：例如{10010}时，这个定点数=1x16+1x2=18

1.2 直通估计器STE

使用STE来规避梯度为0的问题
一个简单的例子是在Bernoulli分布采样中定义STE

其中，c为目标函数，由于Bernoulli采样是一个不可微分的过程，因此$\frac{\partial q}{\partial p}$ 没有定义，因此反向传播不能由链式法则直接求得 $\frac{\partial c}{\partial p}$

然而由于q和p具有相同的期望，可以使用$\frac{\partial c}{\partial p}$ 对$\frac{\partial c}{\partial q}$ 做近似，换句话说STE给出了一个自定义的$\frac{\partial q}{\partial p}$，即$\frac{\partial q}{\partial p}=1$

本文定义$quantiz{e}_k$ 将浮点数输入 $r_i\in [0,1]$ 量化到 k-bit输出 $r_o\in [0,1]$ ，其STE如下：

可以看到STE的输出是一个k-bit实数，而$r_o$是一个k-bit定点数，两个k-bit实数可以使用第一小节描述的定点数点乘（使用正确的缩放）

把这一段原文放在这里，我的理解是，$quantiz{e}_k$ STE的输出是实数，而$quantiz{e}_k$的输出$r_o$为定点数，可以把STE的输出先按照定义的$quantiz{e}_k$转换成定点数，然后再使用比特卷积dot product

1.3 权重的低比特位宽量化

在之前的工作中，STE被用来二值化权重，例如在BNN中

在XNOR-Net中，权重在二值化后进行了缩放

在XNOR-Net中，缩放因子 $E_F(|r_i|)$是对应每个权重的输出通道的绝对值取均值，引入这个scaling factor的可以增加权重值的表达范围，同时仍然能在前向传播卷积时利用比特卷积。

然而在反向传播计算梯度和权重之间的卷积时，channel-wise缩放因子的存在使得无法使用比特卷积核，因此，使用一个常量缩放因子来代替通道级缩放。

当k>1时，使用k-bit表示的权重，将STE应用在权重上：

量化到k-bit之前，使用 $tanh$来限制权重值得范围在-1~1之间，通过将量化的数值限制在[0, 1]之间，最大值是对整个层的权重而言的。然后通过$quantiz{e}_k$将其量化到[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\in [0,1]$，因此在DoReFa-Net中激活的量化简单定义为：

1.5 梯度的低比特位宽量化

在前面几节中描述了确定性量化可以产生低比特位宽的权重和激活，然而发现低位宽梯度必须采用随机量化。

要量化梯度到低比特位宽，需要注意梯度是无界的，并且可能具有比激活值更大的值范围。

其中，$dr=\frac{\partial c}{\partial r}$ 为一些层的输出 r 在反向传播时的梯度（对损失函数c的导数），最大致是对梯度张量的所有维度(除了batch size)的统计，然后在梯度上进行缩放将结果映射到0~1之间，并且在量化之后又缩放回去

为了进一步补偿潜在的由梯度量化所引入的bias，使用一个额外的噪声方程$N(k)=\frac{\sigma }{2^k-1}$，其中$\sigma 服从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))