RepVGG：让VGG风格的ConvNets再次伟大

论文地址：https://arxiv.org/abs/2101.03697

我们提出了一种简单但功能强大的卷积神经网络结构，该模型在推理时类似于VGG，只有3×3的卷积和ReLU堆叠而成，而训练时间模型具有多分支拓扑结构。训练时间和推理时间结构的这种解耦是通过结构重新参数化技术实现的，因此该模型被命名为RepVGG。在ImageNet上，RepVGG达到了超过80%的TOP-1准确率，据我们所知，这是第一次使用普通模型。在NVIDIA 1080Ti GPU上，RepVGG型号的运行速度比ResNet-50快83%，比ResNet-101快101%，精度更高，并且与EfficientNet和RegNet等最先进的型号相比，显示出良好的精度和速度折衷。代码和经过训练的模型可在以下位置获得 https://github.com/megvii-model/RepVGG.

---

1.引言

经典的卷积神经网络(ConvNet) VGG通过一个由conv、ReLU和pooling组成的简单体系结构在图像识别方面取得了巨大成功。随着Inception、ResNet和DenseNet的出现，大量的研究兴趣转移到了精心设计的架构上，使得模型越来越复杂 ，一些最近的架构是基于自动或手动架构搜索，或者搜索基于基本架构的混合尺寸策略得到的强大架构。

虽然许多复杂的卷积网络比简单的卷积网络具有更高的精度，但其缺点是明显的。

1. 复杂的多分支设计（如ResNet中的残差相加和Inception中的分支连接）使模型难以实现和自定义，降低了推理速度和降低了内存利用率。
2. 一些组件（例如Xception和MobileNets中的depth conv和ShuffleNets中的channel shuffle）增加了内存访问成本，缺乏各种设备的支持。

由于影响推理速度的因素太多，浮点运算（FLOPs）的数量并不能精确地反映实际速度。尽管一些新模型的FLOP低于老式模型，如VGG和ResNet-18/34/50，他们可能不会跑得更快，因此，VGG和ResNets的原始版本仍然大量用于学术界和工业界。
在本文中，我们提出了RepVGG，这是一种VGG风格的架构，其性能优于许多复杂的模型（图1）。RepVGG具有以下优点。

• 该模型具有类似VGG的无分支（即前馈）拓扑，这意味着每一层都将其唯一前一层的输出作为输入，并将输出馈送到其唯一后一层。
• 该模型的主体仅使用3×3 conv和ReLU。
• 具体的架构(包括特定的深度和层宽度)实例化时不需要自动搜索、手动细化、复合缩放，也不需要其他繁重的设计。

---

---

对于一个普通模型来说，要达到与多分支体系结构相当的性能水平是很有挑战性的。一种解释是，多分支拓扑，如ResNet，使模型成为众多浅层模型的隐式集成，从而训练多分支模型避免了梯度消失问题。

由于多分支结构的优点都是训练的，而缺点是不利于推理的，我们提出了通过结构重新参数化将训练时多分支结构和推理时平面结构解耦，即通过变换结构参数将结构从一个结构转换到另一个结构。具体地说，网络结构与一组参数相耦合，例如，

卷积层由四阶核张量表示。如果某一结构的参数可以转换为另一结构耦合的另一组参数，我们可以等效地用后者替代前者，从而改变整个网络架构。

具体来说，我们使用identity和1×1分支构造了训练时的RepVGG，这是受ResNet的启发，但采用了不同的方式，可以通过结构重新参数化来删除分支(图2、4)。经过训练后，我们用简单代数进行变换，将一个identity分支看作是一个降级的1×1 conv，后者可以进一步看作是一个降级的3×3 conv，这样我们就可以用原3×3 kernel、identity和1×1分支以及批归一化(BN)层的训练参数构造一个3×3 kernel。因此，转换后的模型有一堆3×3的 conv层，保存用于测试和部署。


值得注意的是，推理时RepVGG的主体只有一种类型的运算符：3x3 conv后跟ReLU，这使得RepVGG在GPU等通用计算设备上速度很快。更好的是，RepVGG允许专门的硬件实现更高的速度，因为考虑到芯片的尺寸和功耗，我们需要的操作员类型越少，我们可以在芯片上集成的计算单元就越多。因此，专门用于RepVGG的推理芯片可以拥有大量的3×3-ReLU单元和更少的存储单元(因为普通的拓扑结构是存储经济的，如图3所示)。我们的贡献总结如下。

• 我们提出了RepVGG，这是一种简单的架构，与最先进的技术相比，具有良好的速度-精度权衡。
• 我们建议使用结构重参数化来解耦一个训练时间的多分支拓扑和一个推断时间的简单架构。
• 我们展示了RepVGG在图像分类和语义分割方面的有效性，以及实现的效率和易用性。

---

2.相关工作

2.1.从单路径到多分支

在VGG将ImageNet分类的top1准确率提高到70%以上之后，为了提高性能，在使ConvNets复杂化方面有了很多创新，如当代的GoogLeNet和后来的Inception模型采用了精心设计的多分支架构，ResNe]提出了简化的双分支架构，DenseNet通过将低层与大量高层连接使拓扑更加复杂。神经体系结构搜索(NAS)和手工设计空间设计可以生成性能更高的ConvNets，但代价是巨大的计算资源或人力资源。一些大版本的NAS生成的模型甚至无法在普通GPU上进行训练，因此限制了应用。除了实现上的不便外，复杂的模型可能会降低并行度，从而减慢推理速度。

2.2.单路径模型的有效训练

已经有一些尝试在没有分支的情况下训练ConvNets。然而，以往的工作主要是试图使非常深的模型以合理的精度收敛，但并不能达到比复杂模型更好的性能。Xiao等人提出了一种初始化方法来训练极深的普通 ConvNet。使用基于平均场理论的方案 (mean-field-theory-based scheme)，10000 层网络在 MNIST 上的训练准确率超过 99%，在 CIFAR-10 上的准确率超过 82%。尽管这些模型并不实用（即使 LeNet-5 在 MNIST 上的准确率可以达到 99.3%，而 VGG-16 在 CIFAR10 上的准确率也可以达到 93% 以上），但其理论贡献是有见地的。最近的一项工作结合了几种技术，包括 Leaky ReLU、最大范数和仔细初始化。在 ImageNet 上，它表明具有 147M 参数的普通 ConvNet 可以达到 74.6% 的 top-1 准确率，比其报告的基线（ResNet-101, 76.6%, 45M 参数）低 2%。

值得注意的是，这篇论文不仅仅是证明普通模型可以相当好地收敛，也不打算训练像ResNet这样的极深的ConvNet。相反，我们的目标是建立一个简单的模型，具有合理的深度和良好的精度-速度权衡，可以用最常见的组件(例如，正则卷积和BN)和简单的代数简单地实现。、

2.3.模型重新参数化

DiracNet 是一种与我们相关的重新参数化方法。它通过将卷积层的内核编码为$\hat{\mathrm{W}}=\mathrm{diag}(\mathbf{a})\mathrm{I}+\mathrm{diag}(\mathbf{b}){\mathrm{W}}_{\text{norm }}$来构建深层平面模型，其中 $\hat{\mathrm{W}}$ 是用于卷积的最终权重（将 4 阶张量视为矩阵），$\mathbf{a}$ 和 $\mathbf{b}$是学习向量，$W_{norm}$ 是归一化的可学习核。与具有相当数量参数的 ResNet 相比，DiracNet 的 top-1 准确率在 CIFAR100 上低 2.29%（78.46% vs. 80.75%），在 ImageNet 上低 0.62%（DiracNet-34 的 72.21% vs. ResNet-的 72.83%） 34)。 DiracNet 在两个方面与我们的方法不同。

• RepVGG的训练时间行为是通过具体结构的实际数据流实现的，之后可以转换为另一个结构，而direcnet只是使用conv内核的另一种数学表达式，以便更容易地优化。换句话说，训练时RepVGG是一个真正的多分支模型，而DiracNet不是。
• DiracNet 的性能高于通常参数化的普通模型，但低于可比较的 ResNet，而 RepVGG 模型的性能大大优于 ResNet。

Asym Conv Block（ACB）、DO Conv 和ExpandNet也可以被视为结构重新参数化，因为它们将块转换为Conv。与我们的方法相比，不同之处在于它们是为组件级改进而设计的，并用作任何架构中Conv层的替代品，而我们的结构重新参数化对于训练普通卷积网至关重要，如第4.2.节所示。

2.4.Winograd卷积

RepVGG仅使用3×3卷积，因为它被一些现代计算库如NVIDIA cuDNN和Intel MKL在GPU和CPU上进行了高度优化。表1显示了在1080Ti GPU上使用cuDNN 7.5.0测试的理论浮点、实际运行时间和计算密度(以每秒 Tera 浮点运算数，TFLOPS衡量)。3×3 conv的理论计算密度约为其他体系结构的4倍，这表明总的理论FLOPs并不能比较不同架构之间的实际速度。Winograd是加速3×3卷积(仅当步长为1时)的经典算法，它得到了cuDNN和MKL等库的良好支持(并在默认情况下启用)。例如，对于标准的F(2×2，3×3)Winograd，3×3卷积的乘法量(MUL)减少到原来的$4/9$。由于乘法比加法更耗时，因此我们计算MUL以测量Winograd支持下的计算成本(在表4,5中用Wino MUs表示)。)。请注意，特定的计算库和硬件决定是否对每个算子使用Winograd，因为由于内存开销，小规模卷积可能不会加速。

---

3.通过结构重新参数构建RepVGG

3.1.简单、快速、节省内存、灵活

使用简单的 ConvNet 至少有三个原因：它们快速、节省内存和灵活。

• Fast 许多最近的多分支架构的理论 FLOP 比 VGG 低，但可能不会运行得更快。例如，VGG-16 的 FLOPs 是 EfficientNet-B3的 8.4 倍，但在 1080Ti 上运行速度快 1.8 倍（表 4），这意味着前者的计算密度是后者的 15 倍。除了 Winograd conv 带来的加速之外，FLOPs 和速度之间的差异可以归因于两个对速度有很大影响但 FLOPs 没有考虑到的重要因素：内存访问成本（MAC）和并行度。例如，尽管所需的分支加法或级联计算可以忽略不计，但 MAC 很重要。此外，MAC 在分组卷积中占时间使用的很大一部分。另一方面，在相同的 FLOP 下，具有高度并行性的模型可能比另一个并行度低的模型快得多。由于多分支拓扑在 Inception 和自动生成的架构中被广泛采用，因此使用多个小型算子而不是几个大型算子。之前的工作报道了 NASNET-A 中的碎片化算子的数量（即一个构建块中的单个conv或pooling操作的数量）为13，这对 GPU 等具有强大并行计算能力的设备不友好并引入了额外的开销，例如内核启动和同步。相比之下，这个数字在 ResNets 中是 2 或 3，我们将其设为 1：单个 conv。

• Memory-economical 多分支拓扑的内存效率低，因为每个分支的结果都需要保存到添加或拼接时，这大大提高了内存占用的峰值。图3显示了残差块的输入需要保持到相加为止。假设块保持特征图大小，额外内存占用的峰值为输入的 2 倍。相反，普通拓扑允许在操作完成时立即释放特定层的输入所占用的内存。在设计专用硬件时，普通的 ConvNet允许深度内存优化并降低内存单元的成本，以便我们可以将更多的计算单元集成到芯片上。

• Flexible 多分支拓扑对架构规范施加了约束。例如，ResNet 要求将卷积层组织为残差块，这限制了灵活性，因为每个残差块的最后一个卷积层必须产生相同形状的张量，否则快捷添加将没有意义。更糟糕的是，多分支拓扑限制了通道修剪的应用，这是一种去除一些不重要通道的实用技术，并且一些方法可以通过自动发现每层的适当宽度来优化模型结构。然而，多分支模型使修剪变得棘手，并导致显着的性能下降或低加速比。相比之下，简单的架构允许我们根据我们的要求自由配置每个conv层并进行修剪以获得更好的性能-效率权衡。

3.2.训练时多分支体系结构

普通模型有许多优点，但其最大的缺点是精度低。论文的结构再参数化基于ResNet，明确构建了一个shortcut分支，对应信息流表示为$y=x+f(x)$ ，通过残差块来学习$f$。当$x$和$f(x)$的维度不匹配时，变成$y=g(x)+f(x)$，其中$g(x)$是以1×1卷积实现的卷积shortcut。ResNets成功的原因可以解释为其多分支的结构使得模型整合了许多浅层模型。当有$n$个模块时，模型可认为整合了$2^n$个模型，因为每个模型有两个分支。

多分支的拓扑结构推断时存在缺点，但是训练时有利，所以只在训练时使用多分支来整合许多模型，为了让多数的成员更浅或者更简单，使用像ResNet的identity(仅当维度匹配时)和1×1分支，使得一个building block的训练信息流为$y=x+g(x)+f(x)$。我们简单的堆叠几个这样的块来构建训练时的模型，当有$n$这样的块时，模型就相当于整合了$3^n$个成员。训练后，该模型被等价地转换到$y=h(x)$，其中$h$通过一个卷积层实现，它的参数是通过一系列代数由训练后的参数导出的。

3.3.简单推理时间模型的再参数化

在本小节中，我们将描述如何将一个经过训练的块转换为一个 $3\times 3$ 的conv层进行推理。注意，在添加之前，我们在每个分支中使用BN(图4)。我们使用 ${\mathrm{W}}^{(3)}\in {\mathbb{R}}^{C_2\times C_1\times 3\times 3}$ 代表一个输入通道为 $C_1$ ，输出通道为 $C_2$ 的 $3\times 3$ 卷积层的卷积核，${\mathrm{W}}^{(1)}\in {\mathbb{R}}^{C_2\times C_1}$ 表示 $1\times 1$ 分支，使用${\mathbf{\mu }}^{(3)},{\mathbf{\sigma }}^{(3)},{\gamma }^{(3)},{\mathbf{\beta }}^{(3)}$作为$3\times 3$卷积后跟的BN层的累计均值、标准方差和学习到的标量因子和bias，${\mathbf{\mu }}^{(1)},{\mathbf{\sigma }}^{(1)},{\gamma }^{(1)},{\mathbf{\beta }}^{(1)}$表示1×1卷积的，${\mathbf{\mu }}^{(0)},{\mathbf{\sigma }}^{(0)},{\gamma }^{(0)},{\mathbf{\beta }}^{(0)}$表示identity分支的，让 ${\mathrm{M}}^{(1)}\in {\mathbb{R}}^{N\times C_1\times H_1\times W_1}$，${\mathrm{M}}^{(2)}\in {\mathbb{R}}^{N\times C_2\times H_2\times W_2}$分别表示输入和输出，$\ast$表示卷积运算。
如果 $C_1=C_2$，$H_1=H_2$，$W_1=W_2$， 则有：

$$\begin{array}{rl}{\mathrm{M}}^{(2)}& =\mathrm{bn}\left({\mathrm{M}}^{(1)}\ast {\mathrm{W}}^{(3)},{\mathbf{\mu }}^{(3)},{\sigma }^{(3)},{\gamma }^{(3)},{\mathbf{\beta }}^{(3)}\right)\\ & +\mathrm{bn}\left({\mathrm{M}}^{(1)}\ast {\mathrm{W}}^{(1)},{\mathbf{\mu }}^{(1)},{\sigma }^{(1)},{\gamma }^{(1)},{\mathbf{\beta }}^{(1)}\right)\\ & +\mathrm{bn}\left({\mathrm{M}}^{(1)},{\mathbf{\mu }}^{(0)},{\sigma }^{(0)},{\gamma }^{(0)},{\mathbf{\beta }}^{(0)}\right).\end{array}$$
否则，则不适用identity分支，上述等式只有两项。这里的bn指的是推断时的BN函数 $\forall 1\le i\le C_2$

$$\mathrm{bn}(\mathrm{M},\mathbf{\mu },\mathbf{\sigma },\gamma ,\mathbf{\beta }{)}_{:,i,:,:}=\left({\mathrm{M}}_{:,i,:,:}-{\mu }_i\right)\frac{{\gamma }_i}{{\mathbf{\sigma }}_i}+{\mathbf{\beta }}_i$$
首先将每个BN和它前面的卷积层转换为一个带有bias向量的卷积，设 { W ′， b ′ } \{ {W′，b′}\} { W′，b′}为核，偏差由 { W ， μ ， σ ， γ ， β } \{ {W，μ，σ，γ，β}\} { W，μ，σ，γ，β}转换，我们得到

$${\mathrm{W}}_{i,\mathrm{i},\mathrm{i},\mathrm{i}}^{\prime }=\frac{{\gamma }_i}{{\sigma }_i}{\mathrm{W}}_{i,\mathrm{i},\mathrm{i},\mathrm{i}},{\mathbf{b}}_i^{\prime }=-\frac{{\mathbf{\mu }}_i{\gamma }_i}{{\mathbf{\sigma }}_i}+{\mathbf{\beta }}_i$$

那么很容易验证 $\forall 1\le i\le C_2$

$$\mathrm{bn}(\mathrm{M}\ast ,\mu ,\sigma ,\gamma ,\beta {)}_{:,\mathrm{i},,:}={\left(\mathrm{M}{\ast }_{\mathrm{W}}^{\prime }\right)}_{:,\mathrm{i},,,:}+{\mathrm{b}}_{\mathrm{i}}^{\prime }$$

以上的变换也应用到了identity分支，因为一个identity映射可以被视作一个以identity矩阵作为卷积核的1×1的卷积。在此变换后，我们将会有1个3×3卷积核，2个1×1卷积核，以及3个bias向量。然后把这3个bias相加得到最终的bias，将1×1卷积核加到3×3卷积核的中心点上得到3×3卷积核，具体实现简单，首先给两个1×1卷积做值为0的padding，扩大到3×3，然后把这3个卷积核相加，如图4所示。注意，这个变换的等价性要求3×3卷积核1×1卷积要有相同stride，并且1×1卷积核的padding配置应该比前者小一个像素。例如3×3的层给输入pad一个像素，则1×1层的padding=0。

---

---

卷积层与BN层融合的公式推导过程：

卷积层公式如下，其中 $W$ 为权重，$b$ 为偏置：
$$Conv(x)=W\ast x+b$$
BN层公式如下，其中 $\gamma$ 与 $\beta$ 为学习参数，$mean$ 为批次样本数据均值，$\sigma$ 为方差，$\epsilon$ 为极小但不为零的数：
$$BN(x)=\gamma \times \frac{x-\text{ mean }}{\sqrt{{\sigma }^2+\epsilon }}+\beta$$

将卷积层结果代入到BN公式中得：
$$\mathrm{BN}(\mathrm{Com}v(x))=\gamma \times \frac{W\ast x+b-\text{ mean }}{\sqrt{{\sigma }^2+\epsilon }}+\beta$$
令 $BN(Conv(x))=y$，进一步化简为：
$$y=\frac{\gamma \times \mathbf{W}\ast \mathbf{x}}{\sqrt{{\sigma }^2+\epsilon }}+\left(\frac{\gamma \times (b-\text{ mean })}{\sqrt{{\sigma }^2+\epsilon }}+\beta \right)$$
又令：
$$\begin{array}{l}{\mathbf{W}}_{\text{fused }}=\frac{\gamma \times \mathbf{W}\ast \mathbf{x}}{\sqrt{{\sigma }^2+\epsilon }}\end{array}$$
$$\begin{array}{l}{b}_{\text{fused }}=\frac{\gamma \times (b-\text{ mean })}{\sqrt{{\sigma }^2+\epsilon }}+\beta \end{array}$$
最终得到：
$$BN(\mathrm{Conv}(x))={\mathbf{W}}_{\text{fused }}\ast x+b_{\text{fused }}$$

---

3.4. 结构规范

表2显示了RepVGG的详细结构。RepVGG采用普通的拓扑结构，大量使用3×3卷积，因为作者希望RepVGG只有一种类型的操作，所以没有像VGG一样采用max pooling。其架构将3×3卷积分为5个stage，第1个stage下采样，stride=2。对于图像分类，使用global average pooling+全连接层作为head，其它特定任务也可以基于任何一层生成的特征图执行。

每个stage的layer数量按照下面三个原则制定：

1. 第一个stage在大的分辨率上执行，比较耗时，所以只用1层来降低延迟;
2. 最后一层一个有更多的通道，所以只用1层来减少参数；
3. 倒数第二层放最多的层（该层有14 × 14输出分辨率），这一点follow了ResNet设计（ResNet-101在其14×14的分辨率stage放置了69层）。

让五个stage所含的layer数量分别为 $1,2,4,14,1$ 则构建了一个名为RepVGG-A的实例模型，$2，3，4$ stage都加 $2$ 层之后得到更深的版本RepVGG-B。RepVGG-A用于和轻权重和中等权重的模型相比，RepVGG-B则和高性能的相比。

每层的宽度，即通道数则在经典宽度设置 $[64,128,256,512]$ 上乘以系数来一致的缩放。$a$ 用来缩放前 $4$ 个stage，$b$用来缩放最后一个stage，通常$b>a$，因为最后一次应该有更丰富的特征用于分类或下游任务。$b$很大也不会显著增加延迟，因为最后一个stage只有一层，由于第一次特征图分辨率很大，所以要求 $a<1$，避免很大的卷积运算，所以第一个stage的 $width$ 是 $min(64,64a)$。

为了进一步减少参数和计算，可以选择用3×3组卷积核密集卷积交替来权衡精度和效率。具体而言，在RepVGG-A的第 $3，5，7，\dots ，21$ 层以及RepVGG-B的第 $23，25，27$ 层设置组卷积，组的数量 $g$为$1，2$ 或者$4$。不适用adjacent groupwise conv。1×1分支应该和3×3卷积一样有相同的 $g$。

---

---

4.实验

4.1. RepVGG在 ImageNet 分类上的表现

4.4.限制

RepVGG是一种简单快速、实用性强的模型，旨在GPU和特定硬件上实现最快速度，不太关注参数或FLOPs数量。尽管RepVGG的参数比ResNet更加高效，但其不如轻量级网络如MobileNets和ShuffleNets等在移动端上的表现。

---

5.结论

论文提出RepVGG，它仅由3×3卷积和ReLU组成，非常适用于GPU或者专用于推理的芯片。在结构化再参数方法的作用下，该简单的卷积网络在ImageNet上的top-1准确率超过了80%，并且相比SOTA模型，在精度和速度上取得更好权衡。

6.代码

# --------------------------------------------------------
# RepVGG: Making VGG-style ConvNets Great Again (https://openaccess.thecvf.com/content/CVPR2021/papers/Ding_RepVGG_Making_VGG-Style_ConvNets_Great_Again_CVPR_2021_paper.pdf)
# Github source: https://github.com/DingXiaoH/RepVGG
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
import torch.nn as nn
import numpy as np
import torch
import copy
from se_block import SEBlock
import torch.utils.checkpoint as checkpoint

def conv_bn(in_channels, out_channels, kernel_size, stride, padding, groups=1):
    result = nn.Sequential()
    result.add_module('conv', nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
                                                  kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, bias=False))
    result.add_module('bn', nn.BatchNorm2d(num_features=out_channels))
    return result

class RepVGGBlock(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size,
                 stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros', deploy=False, use_se=False):
        super(RepVGGBlock, self).__init__()
        self.deploy = deploy
        self.groups = groups
        self.in_channels = in_channels

        assert kernel_size == 3
        assert padding == 1

        padding_11 = padding - kernel_size // 2

        self.nonlinearity = nn.ReLU()

        if use_se:
            #   Note that RepVGG-D2se uses SE before nonlinearity. But RepVGGplus models uses SE after nonlinearity.
            self.se = SEBlock(out_channels, internal_neurons=out_channels // 16)
        else:
            self.se = nn.Identity()

        if deploy:
            self.rbr_reparam = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
                                      padding=padding, dilation=dilation, groups=groups, bias=True, padding_mode=padding_mode)

        else:
            self.rbr_identity = nn.BatchNorm2d(num_features=in_channels) if out_channels == in_channels and stride == 1 else None
            self.rbr_dense = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups)
            self.rbr_1x1 = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=padding_11, groups=groups)
            print('RepVGG Block, identity = ', self.rbr_identity)


    def forward(self, inputs):
        if hasattr(self, 'rbr_reparam'):
            return self.nonlinearity(self.se(self.rbr_reparam(inputs)))

        if self.rbr_identity is None:
            id_out = 0
        else:
            id_out = self.rbr_identity(inputs)

        return self.nonlinearity(self.se(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out))


    #   Optional. This may improve the accuracy and facilitates quantization in some cases.
    #   1.  Cancel the original weight decay on rbr_dense.conv.weight and rbr_1x1.conv.weight.
    #   2.  Use like this.
    #       loss = criterion(....)
    #       for every RepVGGBlock blk:
    #           loss += weight_decay_coefficient * 0.5 * blk.get_cust_L2()
    #       optimizer.zero_grad()
    #       loss.backward()
    def get_custom_L2(self):
        K3 = self.rbr_dense.conv.weight
        K1 = self.rbr_1x1.conv.weight
        t3 = (self.rbr_dense.bn.weight / ((self.rbr_dense.bn.running_var + self.rbr_dense.bn.eps).sqrt())).reshape(-1, 1, 1, 1).detach()
        t1 = (self.rbr_1x1.bn.weight / ((self.rbr_1x1.bn.running_var + self.rbr_1x1.bn.eps).sqrt())).reshape(-1, 1, 1, 1).detach()

        l2_loss_circle = (K3 ** 2).sum() - (K3[:, :, 1:2, 1:2] ** 2).sum()      # The L2 loss of the "circle" of weights in 3x3 kernel. Use regular L2 on them.
        eq_kernel = K3[:, :, 1:2, 1:2] * t3 + K1 * t1                           # The equivalent resultant central point of 3x3 kernel.
        l2_loss_eq_kernel = (eq_kernel ** 2 / (t3 ** 2 + t1 ** 2)).sum()        # Normalize for an L2 coefficient comparable to regular L2.
        return l2_loss_eq_kernel + l2_loss_circle



#   This func derives the equivalent kernel and bias in a DIFFERENTIABLE way.
#   You can get the equivalent kernel and bias at any time and do whatever you want,
    #   for example, apply some penalties or constraints during training, just like you do to the other models.
#   May be useful for quantization or pruning.
    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

    def _pad_1x1_to_3x3_tensor(self, kernel1x1):
        if kernel1x1 is None:
            return 0
        else:
            return torch.nn.functional.pad(kernel1x1, [1,1,1,1])

    def _fuse_bn_tensor(self, branch):
        if branch is None:
            return 0, 0
        if isinstance(branch, nn.Sequential):
            kernel = branch.conv.weight
            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)
            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
            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)
        return kernel * t, beta - running_mean * gamma / std

    def switch_to_deploy(self):
        if hasattr(self, 'rbr_reparam'):
            return
        kernel, bias = self.get_equivalent_kernel_bias()
        self.rbr_reparam = nn.Conv2d(in_channels=self.rbr_dense.conv.in_channels, out_channels=self.rbr_dense.conv.out_channels,
                                     kernel_size=self.rbr_dense.conv.kernel_size, stride=self.rbr_dense.conv.stride,
                                     padding=self.rbr_dense.conv.padding, dilation=self.rbr_dense.conv.dilation, groups=self.rbr_dense.conv.groups, bias=True)
        self.rbr_reparam.weight.data = kernel
        self.rbr_reparam.bias.data = bias
        self.__delattr__('rbr_dense')
        self.__delattr__('rbr_1x1')
        if hasattr(self, 'rbr_identity'):
            self.__delattr__('rbr_identity')
        if hasattr(self, 'id_tensor'):
            self.__delattr__('id_tensor')
        self.deploy = True



class RepVGG(nn.Module):

    def __init__(self, num_blocks, num_classes=1000, width_multiplier=None, override_groups_map=None, deploy=False, use_se=False, use_checkpoint=False):
        super(RepVGG, self).__init__()
        assert len(width_multiplier) == 4
        self.deploy = deploy
        self.override_groups_map = override_groups_map or dict()
        assert 0 not in self.override_groups_map
        self.use_se = use_se
        self.use_checkpoint = use_checkpoint

        self.in_planes = min(64, int(64 * width_multiplier[0]))
        self.stage0 = RepVGGBlock(in_channels=3, out_channels=self.in_planes, kernel_size=3, stride=2, padding=1, deploy=self.deploy, use_se=self.use_se)
        self.cur_layer_idx = 1
        self.stage1 = self._make_stage(int(64 * width_multiplier[0]), num_blocks[0], stride=2)
        self.stage2 = self._make_stage(int(128 * width_multiplier[1]), num_blocks[1], stride=2)
        self.stage3 = self._make_stage(int(256 * width_multiplier[2]), num_blocks[2], stride=2)
        self.stage4 = self._make_stage(int(512 * width_multiplier[3]), num_blocks[3], stride=2)
        self.gap = nn.AdaptiveAvgPool2d(output_size=1)
        self.linear = nn.Linear(int(512 * width_multiplier[3]), num_classes)

    def _make_stage(self, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        blocks = []
        for stride in strides:
            cur_groups = self.override_groups_map.get(self.cur_layer_idx, 1)
            blocks.append(RepVGGBlock(in_channels=self.in_planes, out_channels=planes, kernel_size=3,
                                      stride=stride, padding=1, groups=cur_groups, deploy=self.deploy, use_se=self.use_se))
            self.in_planes = planes
            self.cur_layer_idx += 1
        return nn.ModuleList(blocks)

    def forward(self, x):
        out = self.stage0(x)
        for stage in (self.stage1, self.stage2, self.stage3, self.stage4):
            for block in stage:
                if self.use_checkpoint:
                    out = checkpoint.checkpoint(block, out)
                else:
                    out = block(out)
        out = self.gap(out)
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out


optional_groupwise_layers = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26]
g2_map = {
    
    l: 2 for l in optional_groupwise_layers}
g4_map = {
    
    l: 4 for l in optional_groupwise_layers}

def create_RepVGG_A0(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=1000,
                  width_multiplier=[0.75, 0.75, 0.75, 2.5], override_groups_map=None, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_A1(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=1000,
                  width_multiplier=[1, 1, 1, 2.5], override_groups_map=None, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_A2(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=1000,
                  width_multiplier=[1.5, 1.5, 1.5, 2.75], override_groups_map=None, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_B0(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[1, 1, 1, 2.5], override_groups_map=None, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_B1(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[2, 2, 2, 4], override_groups_map=None, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_B1g2(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[2, 2, 2, 4], override_groups_map=g2_map, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_B1g4(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[2, 2, 2, 4], override_groups_map=g4_map, deploy=deploy, use_checkpoint=use_checkpoint)


def create_RepVGG_B2(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=None, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_B2g2(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=g2_map, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_B2g4(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=g4_map, deploy=deploy, use_checkpoint=use_checkpoint)


def create_RepVGG_B3(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[3, 3, 3, 5], override_groups_map=None, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_B3g2(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[3, 3, 3, 5], override_groups_map=g2_map, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_B3g4(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[3, 3, 3, 5], override_groups_map=g4_map, deploy=deploy, use_checkpoint=use_checkpoint)

def create_RepVGG_D2se(deploy=False, use_checkpoint=False):
    return RepVGG(num_blocks=[8, 14, 24, 1], num_classes=1000,
                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=None, deploy=deploy, use_se=True, use_checkpoint=use_checkpoint)


func_dict = {
    
    
'RepVGG-A0': create_RepVGG_A0,
'RepVGG-A1': create_RepVGG_A1,
'RepVGG-A2': create_RepVGG_A2,
'RepVGG-B0': create_RepVGG_B0,
'RepVGG-B1': create_RepVGG_B1,
'RepVGG-B1g2': create_RepVGG_B1g2,
'RepVGG-B1g4': create_RepVGG_B1g4,
'RepVGG-B2': create_RepVGG_B2,
'RepVGG-B2g2': create_RepVGG_B2g2,
'RepVGG-B2g4': create_RepVGG_B2g4,
'RepVGG-B3': create_RepVGG_B3,
'RepVGG-B3g2': create_RepVGG_B3g2,
'RepVGG-B3g4': create_RepVGG_B3g4,
'RepVGG-D2se': create_RepVGG_D2se,      #   Updated at April 25, 2021. This is not reported in the CVPR paper.
}
def get_RepVGG_func_by_name(name):
    return func_dict[name]



#   Use this for converting a RepVGG model or a bigger model with RepVGG as its component
#   Use like this
#   model = create_RepVGG_A0(deploy=False)
#   train model or load weights
#   repvgg_model_convert(model, save_path='repvgg_deploy.pth')
#   If you want to preserve the original model, call with do_copy=True

#   ====================== for using RepVGG as the backbone of a bigger model, e.g., PSPNet, the pseudo code will be like
#   train_backbone = create_RepVGG_B2(deploy=False)
#   train_backbone.load_state_dict(torch.load('RepVGG-B2-train.pth'))
#   train_pspnet = build_pspnet(backbone=train_backbone)
#   segmentation_train(train_pspnet)
#   deploy_pspnet = repvgg_model_convert(train_pspnet)
#   segmentation_test(deploy_pspnet)
#   =====================   example_pspnet.py shows an example

def repvgg_model_convert(model:torch.nn.Module, save_path=None, do_copy=True):
    if do_copy:
        model = copy.deepcopy(model)
    for module in model.modules():
        if hasattr(module, 'switch_to_deploy'):
            module.switch_to_deploy()
    if save_path is not None:
        torch.save(model.state_dict(), save_path)
    return model