结构重参数化宇宙(Re-parameterization Universe)

0. 前言

一方面，大量研究表明，多分支网络架构的性能普遍优于单分支架构；另一方面，相比多分支架构，单分支架构更有利于部署。那么有没有可能训练时采用多分支架构，而推理时使用单分支呢？我只想说，Re-parameterization is All You Need ！！！

1. Re-parameterization Universe

1.1 RepVGG

为什么要使用VGG式单分支架构：

1. 快：并行度高，单分支架构中的运算是大而整的，而多分支是小而碎的，同样的计算量，“大而整”的运算效率远超“小而碎”

2. 省内存：在add或concat之前，每个分支都占用一份内存来缓存运算结果，add或concat后才能释放
   

3. 硬件友好：VGG式网络架构只有3×3卷积这一种算子，无需额外设计

结构重参数化的核心思想：卷积运算是线性运算，具有可加性。
为了实现结构重参数化，需要特殊设计一种多分支模块，其内部只有线性运算，例如卷积和BN，激活函数需要放到多分支模块外面，如下图所示：

那么如何将多分支转换为单分支呢？下图解释的很清楚，如果能看懂就可以跳过后面的解释。

转换过程可以分为3个步骤：

1. 将1×1卷积和identity转换为3×3卷积：
   1×1卷积通过pad可以直接转换为3×3卷积，这里着重解释一下identity如何转换为3×3卷积。identity的作用是保持输出与输入一致，当输入和输出的通道数均为2时，假设输入分别为$x_1$和$x_2$，输出分别为$y_1$和$y_2$，在不考虑bias的情况下，有：
   $$x_1=y_1=w_1^1\ast x_1+w_1^2\ast x_2$$$$x_2=y_2=w_2^1\ast x_1+w_2^2\ast x_2$$显然，为了使上式成立，$w_1^1$是一个权重为1的1×1卷积，$w_1^2$是一个权重为0的1×1卷积；$w_2^1$是一个权重为0的1×1卷积，$w_2^2$是一个权重为1的1×1卷积。通过填充0就可以转换为3×3卷积，论文中也给出了对应的示意图：
   
   依此类推，若输入和输出通道为$C$，则identity可以表示为下图，其实就是单位矩阵中的每个元素扩展为3×3卷积，文中所提到的单位矩阵就是这个意思。
   

2. 对于每个分支，融合3×3卷积和BN：
   由于每个分支都有个BN层，所以需要将卷积和BN进行融合，也是算子折叠中的常规操作，原理如下：
   卷积可以表示为：$$Conv(x)=W(x)+b$$BN可以表示为：$$BN(x)=\gamma \ast \frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta$$将卷积代入BN中，则有：$$ConvBN(x)=\gamma \ast \frac{W(x)+b-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta =\frac{\gamma \ast W(x)}{\sqrt{{\sigma }^2+ϵ}}+\frac{\gamma \ast (b-\mu )}{\sqrt{{\sigma }^2+ϵ}}+\beta$$
   conv后接BN时，bias通常设置为False，因此，上式可简化为以下形式：
   $$ConvBN(x)=\frac{\gamma \ast W(x)}{\sqrt{{\sigma }^2+ϵ}}+(\beta -\frac{\gamma \ast \mu }{\sqrt{{\sigma }^2+ϵ}})$$
   官方代码如下：

def _fuse_bn_tensor(self, branch):
        if branch is None:
            return 0, 0
        if isinstance(branch, nn.Sequential):
       		# 取出卷积层的权重
            kernel = branch.conv.weight  
            # 取出BN层的参数
            running_mean = branch.bn.running_mean
            running_var = branch.bn.running_var
            gamma = branch.bn.weight
            beta = branch.bn.bias
            eps = branch.bn.eps
        else:
            assert isinstance(branch, nn.BatchNorm2d)
            # 将Identity分支转换为3*3卷积
            if not hasattr(self, 'id_tensor'):
                input_dim = self.in_channels // self.groups
                kernel_value = np.zeros((self.in_channels, input_dim, 3, 3), dtype=np.float32)
                for i in range(self.in_channels):
                    kernel_value[i, i % input_dim, 1, 1] = 1
                self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
            # 获取卷积权重
            kernel = self.id_tensor
            # 获取BN层的参数
            running_mean = branch.running_mean
            running_var = branch.running_var
            gamma = branch.weight
            beta = branch.bias
            eps = branch.eps
        std = (running_var + eps).sqrt()
        t = (gamma / std).reshape(-1, 1, 1, 1)
        # 计算融合卷积和BN后的weight和bias
        return kernel * t, beta - running_mean * gamma / std

3. 融合多分支的weight和bias：

因为每个分支内执行的都是线性运算，因此融合时直接线性求和即可：
$$y=(W_{1}x+b_1)+(W_{2}x+b_2)+(W_{3}x+b_3)=(W_1+W_2+W_3)x+(b_1+b_2+b_3)$$官方代码如下：

def get_equivalent_kernel_bias(self):
        kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)
        kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_1x1)
        kernelid, biasid = self._fuse_bn_tensor(self.rbr_identity)
        return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid

1.2. RepOptimizer

2. 应用

2.1 MobileOne

训练时增加额外的分支数（图中的k Blocks），推理时重参数化为单分支结构，既有多分支的性能，又有单分支的推理速度。

参考资料

[1] RepVGG：极简架构，SOTA性能，让VGG式模型再次伟大（CVPR-2021）
[2] 官方开源代码