网络推理加速的一些基础操作

本页介绍一些遇到的加速的基础操作。

推理时BN层融合入卷积层

参考

To execute neural network inference, kernels are invoked for neural network layers in order to compute the output tensors given the input tensors. Each kernel call will bring some overhead time….To achieve high throughput and low latency for neural network inference, the rule of thumb is to have fewer large kernel calls instead of many small kernel calls.

设卷积核的权重和偏置为\(W\)和\(b\)，卷积层输出和输出分别为\(X\)和\(X^{\prime}\)​，有

\[X^{\prime} = X * W + b\]

卷积核输出\(X^{\prime}\)紧接进入BN层，设BN层的输出为\(Y\)，则有

\[\begin{aligned} Y &= \gamma \frac{X^{\prime} - \mu}{\sqrt{\delta^2 + \epsilon}} + \beta= \gamma \frac{X * W + b - \mu}{\sqrt{\delta^2 + \epsilon}} + \beta \\ &= X * (\frac{\gamma}{\sqrt{\delta^2 + \epsilon}}W) + \beta + \frac{\gamma}{\sqrt{\delta^2 + \epsilon}}(b - \mu) \\ &= X * W^{\prime} + b^{\prime} \end{aligned}\]

即 conv + BN 等效于调用一个权重为\(W^{\prime} = \frac{\gamma}{\sqrt{\delta^2 + \epsilon}}W\)，偏置为\(b^{\prime} = \beta + \frac{\gamma}{\sqrt{\delta^2 + \epsilon}}(b - \mu)\)的新卷积核。该步融合减少了kernel的invoke次数，达到加速的效果。

以下是Pytorch源码中的BN fusion部分

# torch/nn/utils/fusion.py
def fuse_conv_bn_weights(conv_w, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b):
    if conv_b is None:
        conv_b = bn_rm.new_zeros(bn_rm.shape)
    bn_var_rsqrt = torch.rsqrt(bn_rv + bn_eps)

    conv_w = conv_w * (bn_w * bn_var_rsqrt).reshape([-1] + [1] * (len(conv_w.shape) - 1))
    conv_b = (conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b

    return torch.nn.Parameter(conv_w), torch.nn.Parameter(conv_b)

bias correction

高通的一篇工作(DFQ)发现，对网络中的某一层进行权重量化（激活尚未量化）之后，得到的输出相比于量化之前是有偏的。具体来说，对于相同的输入\(x\)，设量化前权重和输出分别为\(W\)和\(y\)，量化后权重和输出分别为\(\widetilde{W}\)和\(\widetilde{y}\)，有

\[E[\widetilde{y}] \neq E[y]\]

其中\(\widetilde{y} = \widetilde{W}x,\quad y = Wx\).

这可能是量化过程中对权重进行clip导致的。而这种输出的分布偏移传到下一层，可能会影响效果，因此需要将这部分输出上的偏移抵消掉。有：

\[\begin{align} E[y] &= E[Wx] \\ &= E[(\widetilde{W} - \epsilon)x] = E[\widetilde{W}x - \epsilon x] \\ &= E[\widetilde{y}] - E[\epsilon x] \\ &= E[\widetilde{y}] - \epsilon E[x] \end{align}\]

其中\(\epsilon = \widetilde{W} - W\)为量化误差。

• 在训练数据（此时无需标签）可得的情况下，具体的做法为，先将一批数据输入全精度模型，得到每一层输出的均值\(E[y_l], \quad l=1, 2, \cdots, L\)，然后将同一批数据通入权重量化后模型，分别统计第\(l\)层的输出期望\(E[\widetilde{y_l}]\)，将偏差\(E[\widetilde{y_l}] - E[y]\)加到该层的bias参数中即可。

• 当训练数据不可得时，DFQ提出可以从该层前的BN层中获取\(E[x]\)，从而得到补偿值\(\epsilon E[x]\)，有兴趣的读者可以阅读原文中详细的描述。

这里高通提出的bias correction是一种后训练（post-training）的方法。参考