ConvNeXt：全面超越 Swin Transformer 的 CNN
date
Mar 10, 2023
Last edited time
Mar 27, 2023 08:28 AM
status
Published
slug
ConvNeXt论文阅读
tags
DL
CV
summary
转载
type
Post
origin
Field
Plat
自从 ViT 提出之后，在过去的一年里（2021 年），基于 transformer 的模型在计算机视觉各个领域全面超越 CNN 模型。然而，这很大程度上都归功于 Local Vision Transformer 模型，Swin Transformer 是其中重要代表。原生的 ViT 模型其计算量与图像大小的平方成正比，而 Local Vision Transformer 模型由于采用 local attention（eg. window attention），其计算量大幅度降低，除此之外，Local Vision Transformer 模型往往也采用金字塔结构，这使得它更容易应用到密集任务如检测和分割中，因为密集任务往往输入图像分辨率较高，而且也需要多尺度特征（eg. FPN）。虽然 Local Vision Transformer 模型超越了 CNN 模型，但是它却越来越像 CNN 了：首先 locality 是卷积所具有的特性，其次金字塔结构也是主流 CNN 模型所采用的设计。那么问题来了，CNN 模型相比 Local Vision Transformer 模型到底差在哪里？能否设计一个纯粹的 CNN 模型实现和 Local Vision Transformer 模型一样的效果？
近日，MetaAI 在论文 A ConvNet for the 2020s 中从 ResNet 出发并借鉴 Swin Transformer 提出了一种新的 CNN 模型：ConvNeXt，其效果无论在图像分类还是检测分割任务上均能超过 Swin Transformer，而且 ConvNeXt 和 vision transformer 一样具有类似的 scalability（随着数据量和模型大小增加，性能同比提升）。
从 ResNet 到 ConvNeXt（1）训练策略（2）宏观设计（3）ResNeXt 化（4）改用 Inverted bottleneck（5）采用 large kernel size（6）微观设计ConvNeXt小结参考
从 ResNet 到 ConvNeXt
论文的整体思路是从原始的 ResNet 出发，通过借鉴 Swin Transformer 的设计来逐步地改进模型。论文共选择了两个不同大小的 ResNet 模型：ResNet50 和 ResNet200，其中 ResNet50 和 Swin-T 有类似的 FLOPs（4G vs 4.5G），而 ResNet200 和 Swin-B 有类似的 FLOPs（15G）。首先做的改进是调整训练策略，然后是模型设计方面的递进优化：宏观设计 ->ResNeXt 化 -> 改用 Inverted bottleneck-> 采用 large kernel size-> 微观设计。由于模型性能和 FLOPs 强相关，所以在优化过程中尽量保持 FLOPs 的稳定。
（1）训练策略
原生的 ViT 需要大规模数据的预训练，而 MetaAI 在 DeiT 论文中提出了一种增强版本的训练策略来解决这个问题，这个训练策略也被随后的 vision transformer 模型所采用。
对于 ResNet50，其训练策略比较简单（torchvision 版本）：batch size 是 32*8，epochs 为 90；优化器采用 momentum=0.9 的 SGD，初始学习速率为 0.1，然后每 30 个 epoch 学习速率衰减为原来的 0.1；正则化只有 L2，weight decay=1e-4；数据增强采用随机缩放裁剪（RandomResizedCrop）+ 水平翻转（RandomHorizontalFlip）。而 DeiT 的训练策略则非常 heavy：采用了比较多的数据增强如 Mixup，Cutmix 和 RandAugment；训练的 epochs 增加至 300；训练的 optimizer 采用 AdamW，学习速率 schedule 采用 cosine decay；采用 smooth label 和 EMA 等优化策略。这里直接将 DeiT 的训练策略（具体参数设置如下表）应用在 ResNet50 模型，其性能从原来的 76.1% 提升至 78.8%（+2.7）。这也说明 vision transformer 模型相比 CNN 模型的提升很多程度上归功于训练策略的优化，关于这个问题，另外一篇论文 Visformer 也论证过这一点，而且最近 timm 库和 torchvison 库也分别发布了 ResNet 新的训练策略，并将 ResNet50 的性能提升至 80%+。
（2）宏观设计
Swin Transfomer 和 CNN 一样也采用金字塔结构：包含 4 个 stage，每个 stage 输出不同尺度的特征。这里考虑 Swin Transformer 和 ResNet 在宏观设计上的区别，主要有两点：每个 stage 的计算量占比以及 stem cell 结构（即 stage1 之前的模块）。 首先是各个 stage 的计算量占比，对比 ResNet50，4 个 stage 的 blocks 数量分别是（3，4，6，3），而 Swin-T 的设置为（2，2，6，2），4 个 stage 的计算量比约为 1:1:3:1。这里调整 ResNet50 各个 stage 的 blocks 数量以和 Swin-T 一致：从原来的（3，4，6，3）调整至（3，3，9，3）。调整之后模型性能从 78.8% 提升至 79.4%，不过这里要注意的一点是，调整后其实 blocks 数量增加了，模型的 FLOPs 从原来的 4G 增加至 4.5G，基本和 Swin-T 一致了，所以这个性能的提升很大程度上归功于 FLOPs 的增加。关于各个 stage 的计算量分配，并没有一个理论上的参考，不过 RegNet 和 EfficientNetV2 论文中都指出，后面的 stages 应该占用更多的计算量。 第二个就是 stem 的区别。对于 Swin-T 模型，其 stem 是一个 patch embedding layer，实际上就是一个 stride=4 的 4x4 conv。而 ResNet50 的 stem 相对更复杂一些：首先是一个 stride=2 的 7x7 conv，然后是一个 stride=2 的 3x3 max pooling。两种 stem 最后均是得到 1/4 大小的特征，所以这里可以直接用 Swin 的 stem 来替换 ResNet 的 stem，这个变动对模型效果影响较小：从 79.4% 提升至 79.5%。对于 ViT 模型，其 patch size 一般较大（eg. 16），只采用一个 stride 较大的 conv 来做 patch embedding 往往会存在一定问题，比如 Mocov3 论文中就指出 patch embed 可能会导致训练不稳定，而论文 Early Convolutions Help Transformers See Better 指出将 patch embed 设计成几个堆叠的 stride=2 的 3*3 conv，无论是在模型效果上，还是在训练稳定性以及收敛速度都更好；而 Swin-T 的 patch size 相对较小，不会出现 ViT 的上述问题，不过 Swin-T 采用的是 non-overlapping conv，后面有论文指出采用 overlapping conv（eg stride=4 的 7x7 conv）会带来一定的性能提升。
（3）ResNeXt 化
相比 ResNet，ResNeXt 通过采用 group conv 来提升性能，标准的 conv 其输入是所有的 channels，而 group conv 会对 channels 进行分组来减少计算量，这样节省下来的计算量用来增加网络的 width 即特征 channels。对于 group conv，其最极端的情况就是每个 channel 一个 group，这样就变成了 depthwise conv（简称 dw conv），dw conv 首先在 MobileNet 中应用，后来也被其它 CNN 模型广泛采用。对于 dw conv，其和 local attention 有很多的相似的地方，local attention 其实就是对 window 里的各个 token 的特征做加权和，而且操作是 per-channel 的；而 dw conv 是对 kernel size 范围的 token 的特征求加权和，也是分 channel 的。这里的最大区别就是：self-attention 的权重是动态计算的（data dependent），而 dw conv 的权重就是学习的 kernel，关于这点，微软的论文 Demystifying Local Vision Transformer: Sparse Connectivity, Weight Sharing, and Dynamic Weight 有深入的分析。
这里将 ResNet50 中的 3x3 conv 替换成 3x3 dw conv，为了弥补 FLOPs 的减少，同时将 ResNet50 的 base width 从原来的 64 增加至 96（和 Swin-T 一致，这里的 base width 是指 stem 后的特征大小），此时模型的 FLOPs 有所增加（5.3G），模型性能提升至 80.5%。
（4）改用 Inverted bottleneck
如果把 self-attention 看成一个 dw conv 的话（这里忽略 self-attention 的 linear projection 操作），那么一个 transformer block 可以近似看成一个 inverted bottleneck，因为 MLP 等效于两个 1x1 conv，并且 MLP 中间隐含层特征是输入特征大小的 4 倍（expansion ratio=4）。inverted bottleneck 最早在 MobileNetV2 中提出，随后的 EfficientNet 也采用了这样的结构。ResNet50 采用的是正常的 residual bottleneck，这里将其改成 inverted bottleneck，即从图（a）变成图（b），虽然 dw conv 的计算量增加了，但是对于包含下采样的 residual block 中，用于 shortcut 的 1x1 conv 计算量却大大降低，最终模型的 FLOPs 减少为 4.6G。这个变动对 ResNet50 的影响较小（80.5%->80.6%)。
（5）采用 large kernel size
自从 VGG 之后，主流的 CNN 往往采用较小的 kernel size，如 3x3 和 5x5，其中 3x3 conv 在 GPU 上有高效的实现。然而 Swin-T 采用的 window size 为 7x7，这比 3x3 conv 对应的 window size 要大得多，所以这里考虑采用更大的 kernel size。 在这之前，首先将 dw conv 移到 inverted bottleneck block 的最开始，如上图（c）所示。对于 transformer block，其实 self-attention 也是在开始，同时由于采用 inverted bottleneck，将 dw conv 移动到最前面可以减少计算量（4.1G），后续采用较大的 kernel size 后模型的 FLOPs 变动更少。由于模型 FLOPs 的降低，模型性能也出现一定的下降：80.6%->79.9%。 然后调整 dw conv 的 kernel size，这里共实验了 5 种 kernel size：3x3，5x5，7x7，9x9 和 11x11。实验发现 kernel size 增加，模型性能有提升，但是在 7x7 之后采用更大的 kernel size 性能达到饱和。所以最终选择 7x7，这样也和 Swin-T 的 window size 一致，由于前面的 dw conv 位置变动，采用 7x7 的 kernel size 基本没带来 FLOPs 的增加。采用 7x7 dw conv 之后，模型的性能又回到 80.6%。
（6）微观设计
经过前面的改动，模型的性能已经提升到 80% 以上，此时改动后的 ResNet50 也和 Swin-T 在整体结构上很类似了，下面我们开始关注一些微观设计上的差异，或者说是 layer 级别的不同。 首先是激活函数，CNN 模型一般采用 ReLU，而 transformer 模型常常采用 GELU，两者的区别如下图所示。这里把激活函数都从 ReLU 改成 GELU，模型效果没有变化（80.6%）。
另外的一个差异是 transformer 模型只在 MLP 中间采用了非线性激活，而 CNN 模型常常每个 conv 之后就会用跟一个非线性激活函数。如下图示，这里只保留中间 1x1 conv 之后的 GELU，就和 Swin-T 基本保持一致了，这个变动使模型性能从 80.6% 提升至 81.3%。
对于 norm 层，也存在和激活函数一样的问题，transformer 中只在 self-attention 和 MLP 的开始采用了 LayerNorm，而 ResNet 每个 conv 之后采用 BatchNorm，比 transformer 多一个 norm 层。这里去掉其它的 BatchNorm，只保留中间 1x1 conv 前的 BatchNorm，此时模型性能有 0.1% 的提升。实际上要和 transformer 保持一致，应该在 block 最开始增加一个 BatchNorm，但是这个并没有提升性能，所以最终只留下了一个 norm 层。另外，transformer 的 norm 层采用 LayerNorm，而 CNN 常采用 BatchNorm，一般情况下 BatchNorm 要比 LayerNorm 效果要好，但是 BatchNorm 受 batch size 的影响较大。这里将 BatchNorm 替换成 LayerNorm 后，模型性能只有微弱的下降（80.5%）。
最后一个差异是下采样，ResNet 中的下采样一般放在每个 stage 的最开始的 block 中，采用 stride=2 的 3x3 conv；但是 Swin-T 采用分离的下采样，即下采样是放在两个 stage 之间，通过一个 stride=2 的 2x2 conv（论文中叫 patch merging layer）。但是实验发现，如果直接改用 Swin-T 的下采样，会出现训练发散问题，解决的办法是在添加几个 norm 层：在 stem 之后，每个下采样层之前以及 global avg pooling 之后都增加一个 LayerNom（Swin-T 也是这样做的）。最终模型的性能提升至 82.0%，超过 Swin-T（81.3%）。
经过 6 个方面的修改或者优化，最终得到的模型称为 ConvNeXt，其模型结构如下所示，可以看到，ConvNeXt-T 和 Swin-T 整体结构基本一致，而且模型 FLOPs 和参数量也基本一样，唯一的差异就是 dw conv 和 W-MSA（MSA 的前面和后面都包含 linear projection，等价 1x1 conv），由于 dw conv 和 W-MSA 的计算量存在差异，所以 ConvNeXt-T 比 Swin-T 包含更多的 blocks。另外 MSA 是 permutation-invariance 的，所以 W-MSA 采用了相对位置编码；而且 Swin-T 需要采用 shifted window 来实现各个 windows 间的信息交互；而相比之下，ConvNeXt-T 不需要这样的特殊处理。
从上面的分析来看，dw conv 的性能是不低于 W-MSA 的，这个其实早在微软的论文 Demystifying Local Vision Transformer: Sparse Connectivity, Weight Sharing, and Dynamic Weight 中被验证过，这篇论文设计了 DW Conv 模型，它也和 Swin 模型结构一致，差别是将 W-MSA 直接换成了 7x7 DW conv，由于保留了两个 linear 层，所以各个 stage 的 blocks 数量也一致。DW Conv-T 和 Swin-T 的性能是一样的（81.3%），略低于 ConvNeXt-T。
ConvNeXt
前面我们从 ResNet50 的 Swin-T 化，最终得到了 ConvNeXt-T 模型，对于更大的模型，可以通过调整特征的维度 C 和各个 stage 的 blocks 数量。论文共设计 5 个 ConvNeXt 模型，其中前 4 个模型分别对标 Swin，而最后 ConvNeXt-XL 是一个更大的模型，用来验证模型的 scalability。
- ConvNeXt-T: C = (96, 192, 384, 768), B = (3, 3, 9, 3)
- ConvNeXt-S: C = (96, 192, 384, 768), B = (3, 3, 27, 3)
- ConvNeXt-B: C = (128, 256, 512, 1024), B = (3, 3, 27, 3)
- ConvNeXt-L: C = (192, 384, 768, 1536), B = (3, 3, 27, 3)
- ConvNeXt-XL: C = (256, 512, 1024, 2048), B = (3, 3, 27, 3)
ConvNeXt 在 ImageNet1K 数据集上的分类效果如下所示，可以看到，ConvNeXt 在不同的 FLOPs 均可以超过 Swin，如果采用 ImageNet21K 预训练后，模型性能有进一步的提升，其中 ConvNeXt-XL 可以达到 87.8%，仅略低于目前的 SOTA 一点（MViT-H, 512^2，88.8%，只有 ImageNet21K）。
在下游任务检测和分割上，ConvNeXt 也可以超过 Swin，如下表所示：
由于 FLOPs 往往不能准确地反映实际的推理速度，所以上面的对比结果也包括模型速度（分类上用 throughout，检测上用 FPS），可以看到在实际推理速度上，ConvNeXt 也略好于 Swin。在显存使用上，两者也相当，比如基于 ConvNeXt-B 的 Cascade Mask-RCNN 需要 17.4GB peak memory，而 Swin-B 需要 18.5GB。另外，如果在 A100 上（支持 TF32），采用 channel last，ConvNeXt 相比 Swin 有明显的速度优势：
在鲁棒性方面，ConvNeXt 也优于 Swin，这说明 ConvNeXt 有很好的泛化。
ConvNeXt 的实现也比较简单，目前已经开源在 https://github.com/facebookresearch/ConvNeXt，这里贴出 Block 的实现，如下所示（采用了 LayerScale）：
小结
从创新上看，ConvNeXt 并没有太多的新意，看起来就是把 vision transformer 一些优化搬到了 CNN 上，而且之前也有很多类似的工作，但我认为 ConvNeXt 是一个很好的工作，因为做了比较全面的实现，而且 ConvNeXt 在工业部署上更有价值，因为 swin transformer 的实现太过 tricky。从 CNN 到 vision transformer 再到 CNN，还包括中间对 MLP 的探索，或许我们越来越能得出结论：架构没有优劣之分，在同样的 FLOPs 下，不同的模型的性能是接近的。但在实际任务上，受限于各种条件，我们可能看到不同模型上的性能差异，这就需要具体问题具体分析了（算法工程师的价值）。