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

最新推荐文章于 2023-01-07 19:07:37 发布

just_sort

最新推荐文章于 2023-01-07 19:07:37 发布

阅读量503

点赞数 1

分类专栏：深度学习论文阅读及算法详解

本文链接：https://blog.csdn.net/just_sort/article/details/104397896

版权

深度学习论文阅读及算法详解专栏收录该内容

102 篇文章 36 订阅

订阅专栏

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

前言

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

原理

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

$Y = W * X + B$

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

$\mu = \frac{1}{m}\sum_{i=1}^mx_i$

$\sigma^2=\frac{1}{m}\sum_{i=1}^m(x_i-\mu)^2$

$x_i=\frac{x_i-\mu}{\sqrt{\sigma^2+ \epsilon}}$

$y_i=\gamma * x_i + \beta$

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

$Y=\gamma*(\frac{(W*X+B)-\mu}{\sqrt{\sigma^2+\epsilon}})+\beta$

$Y=\frac{\gamma*W}{\sqrt{\sigma^2+\epsilon}}*X+\frac{\gamma*(B-\mu)}{\sqrt{\sigma^2+\epsilon}}+\beta$

然后令 $\frac{\gamma}{\sqrt{\delta^2+\epsilon}}$

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

$W_{merged}=W*a$

$B_{merged}=(B-\mu)*a+\beta$

值得一提的是，一般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