深度学习算法优化系列十一 | 折叠Batch Normalization

BBuf

发布于 2020-02-25 12:02:28

9840

发布于 2020-02-25 12:02:28

好久不见，优化系列得继续更啊。。。

前言

今天来介绍一个工程上的常用Trick，即折叠Batch Normalization，也叫作折叠BN。我们知道一般BN是跟在卷积层后面，一般还会接上激活函数，也就是conv+BN+relu这种基本组件，但在部署的时候前向推理框架一般都会自动的将BN和它前面的卷积层折叠在一起，实现高效的前向推理网络，今天我就从原理和实现来讲一讲这个操作。

原理

我们知道卷积层的计算可以表示为：

然后BN层的计算可以表示为：

我们把二者组合一下，公式如下：

然后令

那么，合并BN层后的卷积层的权重和偏置可以表示为：

值得一提的是，一般Conv后面接BN的时候很多情况下是不带Bias的，这个时候上面的公式就会少第二项。

Pytorch代码实现

原理介绍完了，接下来我们来看一下代码实现，就以Pytorch的代码为例好了。

def fuse(conv, bn):
    global i
    i = i + 1
    # ********************BN参数*********************
    mean = bn.running_mean
    var_sqrt = torch.sqrt(bn.running_var + bn.eps)
    gamma = bn.weight
    beta = bn.bias
    # *******************conv参数********************
    w = conv.weight
    if conv.bias is not None:
        b = conv.bias
    else:
        b = mean.new_zeros(mean.shape)

    if(i >= 2 and i <= 7):
        b = b - mean + beta * var_sqrt
    else:
        w = w * (gamma / var_sqrt).reshape([conv.out_channels, 1, 1, 1])
        b = (b - mean)/var_sqrt * gamma + beta
    fused_conv = nn.Conv2d(conv.in_channels,
                         conv.out_channels,
                         conv.kernel_size,
                         conv.stride,
                         conv.padding,
                         groups=conv.groups,
                         bias=True)
    fused_conv.weight = nn.Parameter(w)
    fused_conv.bias = nn.Parameter(b)
    return fused_conv

可以看到这个代码实现和上面介绍的公式完全对应，实现了折叠BN的效果。