A ConvNet for the 2020s

论文链接：https://arxiv.org/pdf/2201.03545.pdf

本篇论文所提出的结构如上图。

上图是对标准的ConvNet (ResNet)进行了现代化的过程，包括1)宏观设计，2)ResNeXt, 3)倒置瓶颈，4)大内核大小，5)各种层次的微观设计，设计了层次视觉转换器(Swin)，而没有引入任何基于注意力的模块。前景条是ResNet-50/Swin-T flops模式的模型精度;ResNet-200/ Swin-B状态的结果显示为灰色条形。两种方法的详细结果在后文中会提出。许多Transformer体系结构选择可以合并到ConvNet中，它们会带来越来越好的性能。最后，纯粹的ConvNet模型，命名为ConvNeXt，可以胜过Swin Transformer。

一、训练过程优化
在本论文中，第一步是用vision Transformer训练程序训练一个baseline model ，在本例中是ResNet50/200。在作者的研究中，使用的训练方法接近DeiT和Swin Transformer的。训练从ResNets最初的90个epoch扩展到300个epoch。使用AdamW优化器、数据增强技术(如Mixup、Cutmix、RandAugment、Random Erasing)和正则化方案(包括Random Depth和Label Smoothing)。对于baseline model本身而言，这种增强训练方法将ResNet-50模型的性能从76.1%提高到78.8%(+2.7%)，这意味着传统卷积神经网络和vision Transformer之间的性能差异很大一部分可能是由于训练技术。在整个“现代化”过程中，将使用这个固定的训练方案，并使用相同的超参数。ResNet-50方法的每一个报告的准确性都是由三种不同的随机种子训练获得的平均值。
二、宏观设计
（一）、改变阶段计算比率
ResNet中跨阶段的计算分布最初的设计很大程度上是基于经验的。沉重的“res4”阶段是为了与下游的任务兼容，比如对象检测，其中一个检测器头在14×14特征平面上操作。另一方面，Swin-T遵循相同的原理，但阶段计算比例略有不同，为1:1:3:1。对于更大的Swin Transformer，比例是1:1:9:1。按照这种设计，作者将每个阶段的块数从ResNet-50中的(3,4,6,3)调整为(3,3,9,3)，这也将FLOPs与Swin-T对齐。这将模型的精度从78.8%提高到79.4%。值得注意的是，研究人员已经对计算的分布进行了深入的研究，可能存在更优化的设计。
（二）、改变stem的“Patchify”
通常情况下，stem的设计关注的是在网络开始时如何处理输入的图像。由于自然图像固有的冗余性，在标准卷积神经网络(ConvNets)和视觉变形器(vision transformer)中，普通stem会积极地将输入图像采样到合适的特征图大小。在标准ResNet中，stem包含一个stride 为2的7×7卷积层，然后是一个最大池，其结果是输入图像的4倍下采样。在vision transformer中，使用了一种更激进的“patchify”策略作为stem cell，它对应于较大的核大小(例如核大小= 14或16)和非重叠卷积。Swin Transformer使用了类似的“patchify”层，但patch size的大小更小，为4，以适应体系结构的多阶段设计。我作者将resnet类型的stem cell替换为一个patchify层，该层使用4×4, stride为4卷积层实现。准确率从79.4%提高到79.5%。这表明ResNet中的stem cell可以被一个更简单的“patchify”层取代。
三、Inverted Bottleneck
作者尝试采用ResNeXt的思想，它比普通的ResNet具有更好的FLOPs/准确性权衡。核心部分是分组卷积，其中卷积被分成不同的组。在高层次上，ResNeXt的指导原则是“使用更多的组，扩大宽度”。更准确地说，ResNeXt对瓶颈块中的3×3 conv层使用了分组卷积。因此，可以通过扩大网络宽度来弥补容量损失。在本文中，作者使用深度卷积，这是分组卷积的一种特殊情况，分组的数量等于通道的数量。Depthwise conv已经由MobileNet[32]和Xception推广。深度卷积类似于自注意中的加权和操作，其操作基于每个通道，即只在空间维度上混合信息。深度卷积的使用有效地减少了网络的FLOPs，达到了和预期的一样精度。按照ResNeXt中提出的策略，我们将网络宽度增加到与Swin-T相同的通道数量(从64个增加到96个)。这使得网络性能达到80.5%，FLOPs增加(5.3G)。

每一个Transformer block重要的设计是它创建一个反向瓶颈，如上图，向MLP的隐藏维度块的宽度是输入维度的四倍。这种Transformer设计与在ConvNets中使用的扩展比为4的反向瓶颈设计有关。这个想法是由MobileNetV2推广的，随后在几种先进的ConvNet架构中被引用。

上图中的。(a)是ResNeXt块;(b)是创建的一个反向的瓶颈块
尽管深度卷积层的FLOPs增加了，但这一变化使整个网络FLOPs减少到4.6G，因为下采样残差块的跳连接的1×1 conv层的FLOPs显著减少。这会略微提高性能(从80.5%提高到80.6%)。在ResNet-200 / Swin-B模式中，这一步带来了更多的收益(81.9%至82.6%)，同时也减少了FLOPs。

四、Large Kernel Sizes

(b)inverted bottleneck block，在©空间深度转换层的位置向上移动
(一)、上移动深度转换层
要探索大型内核，一个先决条件是向上移动深度转换层的位置（上图中的 (b)到©）。这是一个设计决策，在transformer中也很明显:MSA块被放置在MLP层之前。由于之前已经得到了一个inverted bottleneck block，这是一个自然的设计选择——复杂/低效的模块(MSA，大内核conv)将有更少的通道，而高效、密集的1×1层将完成繁重的工作。这个中间步骤将FLOPs降低到4.1G，导致性能暂时下降到79.9%。
（二）、增加卷积核大小

作者实验了几种内核大小，包括3、5、7、9和11。网络的性能从79.9% (3×3)提高到80.6% (7×7)，而网络的FLOPs致保持不变。此外，作者观察到较大内核大小的好处在7×7达到饱和点。也在大容量模型中验证了这一行为:当作者将内核大小增加到7×7以上时，ResNet-200机制模型不会显示出进一步的增益。故将在每个块中使用7×7 depth - wise conv。具体结构如上
五、用GELU取代ReLU
NLP和视觉架构之间的一个差异是使用哪种激活函数的细节。随着时间的推移，人们开发了许多激活函数，但Rectified Linear Unit(ReLU)由于其简单和高效，仍然广泛应用于卷积神经网络。在原始的Transformer论文中，ReLU也被用作激活函数。Rectified Linear Unit（GELU），可以被认为是ReLU的平滑变体，在最先进的Transformers中被使用，包括谷歌的BERT和OpenAI的GPT-2，以及最近的ViTs。在ConvNet中，ReLU也可以被GELU替代，但精度没有变化(80.6%)。
六、更少的激活函数
上
Transformer和ResNet块之间的一个小区别是Transformer有更少的激活函数。考虑一个具有key/query/value的线性嵌入层、投影层和MLP块中的两个线性层的Transformer块。MLP块中只有一个激活函数。相比之下，通常的做法是在每个卷积层上附加一个激活函数，包括1 × 1卷积。在这里，作者将研究当坚持相同的策略时，性能是如何变化的。如图4所示，作者从残差块中消除了所有GELU层，除了两个1 × 1层之间的一个，复制了Transformer块的样式。该操作将结果提高了0.7%，达到81.3%，与Swin-T的性能相当。
七、更少的归一化层
Transformer blocks通常也有较少的Transformer blocks。这里作者去掉了两个BatchNorm (BN)层，在conv 1 × 1层之前只留下一个BN层。这进一步提升到了81.4%，已经超过了 Swin-T的结果。每个块的BatchNorm (BN)层甚至比transformer还要少，作者发现在块的开始添加一个额外的BN层并不能改善性能。
八、用LN代替BN
BatchNorm是卷积神经网络的重要组成部分，因为它提高了收敛性并减少了过拟合。然而，BN也有许多错综复杂的地方，会对模型的性能产生不利影响。人们曾多次尝试开发替代的标准化技术，但BN仍然是大多数视觉任务的首选方法。另一方面，在transformer中使用了更简单的Layer Normalization(LN)，从而在不同的应用场景中获得了良好的性能。直接用LN代替原来ResNet中的BN会导致性能不佳。随着网络架构和训练技术的所有修改，继续探索使用LN代替BN的影响。作者观察到卷积神经网络模型在LN的训练中没有任何困难;实际上，该算法的性能略好，达到了81.5%的精度。
九、分离下采样层
在ResNet中，空间下采样由每个阶段开始时的残差块实现，使用stride为2的3×3 conv (在跳连接处使用stride为2的1×1 conv )。但在Swin transformer中，在各个阶段之间添加了一个单独的下采样层。使用stride为2的2×2 conv层和进行空间下采样。这种改变会导致不同的训练。进一步的研究表明，在空间分辨率发生变化的地方增加归一化层有助于稳定训练。这包括几个也在Swin transformer中使用的LN层:一个在每个下采样层之前，一个在主干之后，一个在最终的全局平均池之后。可以将准确率提高到82.0%，显著超过Swin-T的81.3%。
十、实验结果

作者构造了不同的ConvNeXt变体，即ConvNeXtT/S/B/L，其复杂性与 Swin-T/S/B/L相似。ConvNeXt-T/B分别是ResNet-50/200体系“现代化”的最终模型。此外，作者还构建了一个更大的ConvNeXt- XL以进一步测试ConvNeXt的可扩展性。变体只在通道C的数量和每个阶段的块B的数量上有所不同。在ResNets和Swin transformer之后，通道的数量在每个新阶段都翻了一番。配置总结如上。

（一）、在ImageNet-1K中
在精度-算力权衡方面，以及推理吞吐量方面，ConvNeXt与两个强ConvNetbaselines(RegNet和EfficientNet)具有良好的竞争优势。在类似的复杂性下，ConvNeXt的表现也优于Swin Transformer，有时甚至有很大的提升(例如，ConvNeXt- T的提升为0.8%)。没有诸如移动窗口或相对位置偏差等特殊模块，与Swin transformer相比，ConvNeXts有更高的吞吐量。其中一个亮点是在输入图像分辨率为$384^2$的ConvNeXt-B:它的表现比Swin-B高出0.6%(85.1%到84.5%)推理的吞吐量提高12.5%(95.7到85.1张/秒)。当分辨率从$224^2$提高到$384^2$时，ConvNeXt-B相对于Swin-B的FLOPs/吞吐量优势变得更大。此外，当进一步扩展到ConvNeXt-L时，我们观察到85.5%的提升结果。

（二）、在ImageNet-22K中
结果在上表(下)中展示了由ImageNet-22K预训练调优的模型的结果。结果表明在使用大数据集进行预训练时，并不逊色于vision transformer——ConvNeXts仍然与同样大小的Swin transformer表现相同或更好，具有略高的吞吐量。此外，ConvNeXt-XL模型达到了87.8%的精度——相对于$384^2$的ConvNeXt-L来说是一个不错的改进，这表明convnext是可扩展的架构。
在本次消融实验中，作者验证了ConvNeXt块设计是否可归纳到ViT风格的相同的架构，这种结构没有下采样层，并且在所有深度保持相同的特征分辨率(例如14×14)。使用与ViT-S /B/L(384/768/1024)相同的特征维来构造各向同性的ConvNeXt-S/B/L。深度设置为18/18/36，以匹配参数和FLOPs的数量。块结构保持不变。对于ViT-S /B，作者使用DeiT的监督训练结果，对于ViT-L，作者使用MAE的监督训练结果，因为他们采用了比原始ViT改进的训练程序。ConvNeXt模型的训练与之前的设置相同，但热身时间更长。2242分辨率的ImageNet-1K的结果如下表所示。ConvNeXt的性能通常与ViT相当，这表明ConvNeXt块设计在非分层模型中具有竞争力。

（三）、基于COCO的目标检测与分割

上表显示了Swin Transformer、ConvNeXt和ResNeXt等传统ConvNet的对象检测和实例分割结果。在不同的模型复杂性中，ConvNeXt实现了与Swin Transformer相同或更好的性能。当放大到更大的模型(ConvNeXt- b /L/XL)预先训练的ImageNet-22K，在许多情况下，在box and mask AP方面ConvNeXt比Swin transformer要好。
（四）、模型效率
相同的的FLOPs下，与只有密集卷积的ConvNets相比，具有深度卷积的模型更慢，消耗更多内存。人们很自然地会问，ConvNeXt的设计是否会使它实际上效率低下。正如整篇文章所证明的，ConvNeXts的推理吞吐量与Swin transformer相当或超过。这对于分类和其他需要更高分辨率输入的任务来说都是正确的。训练ConvNeXts比训练Swin transformer需要更少的内存。例如，使用ConvNeXt-B骨干训练Cascade Mask-RCNN消耗17.4GB的峰值内存，每个gpu的批大小为2，而swin - b的参考编号是18.5GB。与普通的ViT相比，由于本地计算，ConvNeXt和Swin Transformer都表现出更有利的准确性- flops权衡。值得注意的是，这种效率的提高是由于与视觉Transformers中的自我注意机制没有直接关系。
十一、实验中的超参数设置

上表为 ImageNet-1K/22K的预训练设置。每个模型(例如，ConvNeXt-T/S/B/L)分别有多个随机depth rates(例如，0.1/0.4/0.5/0.5)。

在ImageNet-1K微调结构的超参数设置，每个模型(例如，ConvNeXt-B/L)分别有多个值(例如，0.8/0.95)。调从训练前获得的最终模型权值开始，不使用EMA权值，没有观察到改善。唯一的例外是在ImageNet-1K上进行预训练的convext - L，由于过拟合的原因，其模型精度明显低于EMA的精度，在预训练时选择其最好的EMA模型作为调优的起点。使用分层学习速率衰减，每3个连续的块组成一个组。当模型在$384^2$分辨率下进行微调时，则在随后的测试中使用1.0的裁剪比(即不裁剪)，而不是在$224^2$时进行0.875的裁剪。
十二、框架实现过程中的具体结果
（一）、鲁棒性评价

ConvNeXt(特别是大规模模型变体)展示了有前景的鲁棒性行为，在几个基准测试中优于最先进的 robust transformer。使用额外的ImageNet- 22k数据，ConvNeXtXL展示了强大的域泛化能力(例如，在ImageNet-A/R/Sketch上实现的准确率分别为69.3%/68.2%/55.0%)。这些稳健性评估结果是在没有使用任何专门模块或额外的微调程序的情况下获得的。
（二）、现代化ResNets过程的详细结果

本文提供了ResNet-50 / Swin-T和ResNet-50 / Swin-T两种现代化实验过程的详细结果表。在ImageNet-1K上每一步的top-1精度和FLOPs如上表所示。ResNet-50架构采用3种随机种子。对于ResNet-200，每个阶段的初始块数为(3,24,36,3)，在改变阶段比的阶段，作者将其改为Swin-B的(3,3,27,3)。这大大减少了FLOPs，同时，也将宽度从64增加到84，以保持FLOPs在一个相似的水平。在采用深度卷积的步骤之后，、进一步将宽度增加到128(与Swin-B的相同)作为一个单独的步骤。ResNet-200的观测结果与前文中描述的ResNet-50的观测结果基本一致。一个有趣的区别是，相对于ResNet-50，反向维度带来了ResNet-200体系更大的改进(+0.79% vs +0.14%)。增加内核大小所获得的性能似乎也在内核大小为5而不是7时达到饱和。与ResNet-50模式相比，使用更少的规范化层也有更大的收益(+0.46% vs +0.14%)。

（三）详细的体系结构

上表中给出了ResNet-50、ConvNeXt-T和Swin-T的详细架构的比较。对于不同大小的convnext，只有每个阶段的块数和通道数与ConvNeXt-T不同。ConvNeXts继承ConvNets的简单性，但在视觉识别方面与Swin transformer竞争。
（四）、A100 GPUs基准测试

在之前的实验中作者已经证明，在参数数量相似的情况下，ConvNeXt的推理速度比Swin Transformer略快。现在在更先进的A100 GPUs上对它们进行基准测试，A100 GPUs支持TensorFloat32 (TF32)A100 GPUs。使用PyTorch 1.10中的“Channel Last”内存布局来进一步加速。结果如上表。Swin transformer和ConvNeXts都实现了比A100 GPUs更快的推理吞吐量，但现在ConvNeXts的优势明显更大，有时可以快49%。这项初步研究显示了ConvNeXt应用于标准的ConvNet模块和简单的设计，可以在现代硬件上实际更有效的模型。
十三、总结
ConvNeXt可能更适合某些任务，而transformer可能更适合其他任务。一个恰当的例子是多模态学习，在这种学习中，交叉注意模块可能更适合于建模多个模态之间的特征交互。此外，当用于需要离散、稀疏或结构化输出的任务时，transformer可能更灵活。架构的选择应该满足当前任务的需要，同时力求简单。
最后附上文中架构的实现代码

import tensorflow as tf
import tensorflow_addons as tfa
import tensorflow.keras.layers as layers

# 初始化参数
kernel_initial = tf.keras.initializers.TruncatedNormal(stddev=0.2)
bias_initial = tf.keras.initializers.Constant(value=0)

# 定义一个2*2conv进行下采样
class Downsampling(tf.keras.Sequential):
    def __init__(self, out_dim):
        super(Downsampling, self).__init__([
            layers.LayerNormalization(),
            layers.Conv2D(
                out_dim, kernel_size=2, strides=2, padding='same',
                kernel_initializer=kernel_initial, bias_initializer=bias_initial
            )
        ])

# 定义一个ConvNextBlock(
class ConvNextBlock(layers.Layer):
    """ConvNeXt Block using implementation 1
    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back

    Args:
        dim (int): Number of input channels.
        drop_path (float): Stochastic depth rate. Default: 0.0
        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
    """

    def __init__(self, dim, drop_prob=0, layer_scale_init_value=1e-6):
        super().__init__()
        self.layers = tf.keras.Sequential([
            layers.Conv2D(dim, kernel_size=7, padding="same", groups=dim,
                          kernel_initializer=kernel_initial, bias_initializer=bias_initial),
            layers.LayerNormalization(),
            layers.Conv2D(dim, kernel_size=1, padding="valid",
                          kernel_initializer=kernel_initial, bias_initializer=bias_initial),
            layers.Activation('gelu'),
            layers.Conv2D(dim, kernel_size=1, padding="valid",
                          kernel_initializer=kernel_initial, bias_initializer=bias_initial),
        ])
#设置图层进行缩放的初始值，如果不设置则不进行缩放
        if layer_scale_init_value > 0:
            self.layer_scale_gamma = tf.Variable(
                initial_value=layer_scale_init_value*tf.ones((dim)))
        else:
            self.layer_scale_gamma = None
        self.stochastic_depth = tfa.layers.StochasticDepth(drop_prob)

    def call(self, x):
        skip = x
        x = self.layers(x)
        if self.layer_scale_gamma is not None:
            x = x * self.layer_scale_gamma
        x = self.stochastic_depth([skip, x])
        return x


class ConvNext(tf.keras.Model):
    """ A Tensorflow impl of : `A ConvNet for the 2020s`
        https://arxiv.org/pdf/2201.03545.pdf
    Args:
        in_chans (int): Number of input image channels. Default: 3
        num_classes (int): Number of classes for classification head. Default: 1000
        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
        drop_path_rate (float): Stochastic depth rate. Default: 0.
        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
    """

    def __init__(self, num_classes=1000,
                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], drop_path_rate=0.,
                 layer_scale_init_value=1e-6):
        super().__init__()
        self.downsample_layers = []
        self.downsample_layers.append(tf.keras.Sequential([
            layers.Conv2D(dims[0], kernel_size=4, strides=4, padding="valid"),
            layers.LayerNormalization()
        ]))
        self.downsample_layers += [Downsampling(dim) for dim in dims[1:]]
        self.convnext_blocks = [tf.keras.Sequential([ConvNextBlock(dim, drop_path_rate, layer_scale_init_value) for _ in range(
            depths[i])]) for i, dim in enumerate(dims)]
        self.head = layers.Dense(
            num_classes, kernel_initializer=kernel_initial, bias_initializer=bias_initial)
        self.gap = layers.GlobalAveragePooling2D()
        self.norm = layers.LayerNormalization()

    def call_features(self, x):
        for i in range(4):
            x = self.downsample_layers[i](x)
            x = self.convnext_blocks[i](x)
        x = self.gap(x)
        return self.norm(x)

    def call(self, x):
        x = self.call_features(x)
        x = self.head(x)
        return x