[人工智能]理解ConvNeXt网络（结合代码）

目录

?1.简介

?2.ConvNeXt的设计与实验

2.1 macro design（大的结构上的设计）

2.1.1 Changing stage compute ratio（改变每个stage的堆叠次数）

2.1.2 Changing stem to “Patchify”（stem为最初的下采样模块，改为与swin相似的patch卷积进行下采样）

2.2 ResNeXt（参考ResNeXt）

2.2.1 depth conv（普通卷积改为DW 卷积）

2.2.2 width（增加每个stage网络的深度）

2.3 inverted bottleneck

2.3.1 inverting dims（bottleneck由Resnet中两头粗中间细改为了类似Mobilenetv2中两头细中间粗的结构）

2.4 large kerner size（增加卷积核大小）

2.4.1 Moving up depthwise conv layer（将bottleneck中DW卷积模块上移）

2.4.2 Increasing the kernel size（增大DW卷积的卷积核大小）

2.5 various layer-wise micro designs（微小结构的设计）

2.5.1 Replacing ReLU with GELU（激活函数的改变）

2.5.2?Fewer activation functions（使用更少的激活函数，仅在DW卷积后的全连接层使用激活函数）

?编辑

2.5.3?Fewer normalization layers（将BN替换为LN，并且只在DW卷积后使用LN）

2.5.4?Substituting BN with LN（将BN替换为LN）

2.5.5?Separate downsampling layers（将Resnet中在bottleneck中进行下采样改为使用单独的下采样层）

?3.?ConvNeXt 的版本

4.ConvNeXt的整体结构

?5.?ConvNeXt网络模型代码

1.简介

ConvNeXt网络本身没有什么亮点，全是应用的现有的方法来进行网络的调整，特别是大量细节的设计都是参考了swin transformer的网络结构的。并且ConvNeXt是以ResNet50网络为backbone来进行调整的，所以ConvNeXt的网络结构非常简单，一目了然，理解起来也是非常容易的。并且不仅精度比swin Transformer高，推理速度还快。

综合来说，ConvNeXt是一个非常好的文章。这里放上我看到的一个网友对ConvNeXt网络的评价。

“感觉这篇论文的思路是照着swin-transformer的结构靠拢的，可以明显的看出来很多设计都是损于推理速度的，比如激活函数的选择，组卷积的使用，分支数量的增加，这是一篇很不错的论文，其给出了详细的设计思路同时又留下了很多针对于推理速度的改进空间，我觉得现在卷积神经网络的另一个方向可以研究等效，即repvgg论文中使用的，训练时采用残差结构提高精度，而推理时转化为单分支结构极大的提高推理速度，因为各大ai芯片以及nvidia对3*3卷积的优化的非常的好，以我个人浅薄的认知而言，卷积神经网络的瓶颈还远远没有达到，未来在推理速度的研究方面，卷积神经网络因为其简单的结构很有可能会再次走在transformer的前面。”

?2.ConvNeXt的设计与实验

作者以通过ViT的训练策略训练的Resnet50网络（精度78.8）作为基准网络进行调整，最后能达到82.0的准确率（高于swin-T的81.3），说明将swin Transformer的结构和训练策略应用到Resnet上是很有效果的。

ConvNeX网络的改进主要有以下五个方面：

• 2.1 macro design（大的结构上的设计）

  • 2.1.1 Changing stage compute ratio（改变每个stage的堆叠次数）

    在原ResNet网络中，一般conv4_x（即stage3）堆叠的block的次数是最多的。ResNet50中stage1到stage4堆叠block的次数是(3, 4, 6, 3)比例大概是1:1:2:1，但在Swin Transformer中，比如Swin-T的比例是1:1:3:1，Swin-L的比例是1:1:9:1。很明显，在Swin Transformer中，stage3堆叠block的占比更高。所以作者就将ResNet50中的stage中的堆叠次数由(3, 4, 6, 3)调整成(3, 3, 9, 3)，和Swin-T拥有相似的FLOPs。进行调整后，准确率由78.8%提升到了79.4%。

    

  • 2.1.2 Changing stem to “Patchify”（stem为最初的下采样模块，改为与swin相似的patch卷积进行下采样）

    在之前的卷积神经网络中，一般最初的下采样模块stem一般都是通过一个卷积核大小为7x7步距为2的卷积层以及一个步距为2的最大池化下采样共同组成，高和宽都下采样4倍。但在Transformer模型中一般都是通过一个卷积核非常大且相邻窗口之间没有重叠的（即stride等于kernel_size）卷积层进行下采样。比如在Swin Transformer中采用的是一个卷积核大小为4x4步距为4的卷积层构成patchify，同样是下采样4倍。所以作者将ResNet中的stem也换成了和Swin Transformer一样的patchify。替换后准确率从79.4% 提升到79.5%，并且FLOPs也降低了一点。

• 2.2 ResNeXt（参考ResNeXt）

  • 2.2.1 depth conv（普通卷积改为DW 卷积）

    借鉴了ResNeXt中的组卷积grouped convolution，因为ResNeXt相比普通的ResNet而言在FLOPs以及accuracy之间做到了更好的平衡。而作者采用的是更激进的depthwise convolution，即group数和通道数channel相同。这样做的另一个原因是作者认为depthwise convolution和self-attention中的加权求和操作很相似。

  • 2.2.2 width（增加每个stage网络的深度）

    将最初的通道数由64调整成96和Swin Transformer保持一致，最终准确率达到了80.5%。

• 2.3 inverted bottleneck

  • 2.3.1 inverting dims（bottleneck由Resnet中两头粗中间细改为了类似Mobilenetv2中两头细中间粗的结构）

    作者认为Transformer block中的MLP模块非常像MobileNetV2中的Inverted Bottleneck模块，即两头细中间粗。作者采用Inverted Bottleneck模块后，在较小的模型上准确率由80.5%提升到了80.6%，在较大的模型上准确率由81.9%提升到82.6%。

• 2.4 large kerner size（增加卷积核大小）

  • 2.4.1 Moving up depthwise conv layer（将bottleneck中DW卷积模块上移）

    将depthwise conv模块上移，原来是1x1 conv -> depthwise conv -> 1x1 conv，现在变成了depthwise conv -> 1x1 conv -> 1x1 conv。这么做是因为在Transformer中，MSA模块是放在MLP模块之前的，所以这里进行效仿，将depthwise conv上移。这样改动后，准确率虽然下降到了79.9%，但同时FLOPs也减小了。

  • 2.4.2 Increasing the kernel size（增大DW卷积的卷积核大小）

    作者将depthwise conv的卷积核大小由3x3改成了7x7（和Swin Transformer一样），当然作者也尝试了其他尺寸，包括3, 5, 7, 9, 11发现取到7时准确率就达到了饱和。并且准确率从79.9% (3×3)?增长到?80.6% (7×7)。

• 2.5 various layer-wise micro designs（微小结构的设计）

  • 2.5.1 Replacing ReLU with GELU（激活函数的改变）

    在Transformer中激活函数基本用的都是GELU，而在卷积神经网络中最常用的是ReLU，于是作者又将激活函数替换成了GELU，替换后发现准确率没变化。

  • 2.5.2?Fewer activation functions（使用更少的激活函数，仅在DW卷积后的全连接层使用激活函数）

    使用更少的激活函数。在卷积神经网络中，一般会在每个卷积层或全连接后都接上一个激活函数。但在Transformer中并不是每个模块后都跟有激活函数，比如MLP中只有第一个全连接层后跟了GELU激活函数。接着作者在ConvNeXt Block中也减少激活函数的使用，如下图所示，减少后发现准确率从80.6%增长到81.3%。 不太理解这里的ConvNeXt为什么不添加SE注意力模块，是因为在swin Transformer中没有对应的结构吗。

  • 

  • 2.5.3?Fewer normalization layers（将BN替换为LN，并且只在DW卷积后使用LN）

    在Transformer中，Normalization使用的也比较少，接着作者也减少了ConvNeXt Block中的Normalization层，只保留了depthwise conv后的Normalization层。此时准确率已经达到了81.4%，已经超过了Swin-T。

  • 2.5.4?Substituting BN with LN（将BN替换为LN）

    在Transformer中基本都用的Layer Normalization（LN），因为最开始Transformer是应用在NLP领域的，BN又不适用于NLP相关任务。接着作者将BN全部替换成了LN，发现准确率还有小幅提升达到了81.5%。

  • 2.5.5?Separate downsampling layers（将Resnet中在bottleneck中进行下采样改为使用单独的下采样层）

    在ResNet网络中stage2-stage4的下采样都是通过将主分支上3x3的卷积层步距设置成2，捷径分支上1x1的卷积层步距设置成2进行下采样的。但在Swin Transformer中是通过一个单独的Patch Merging实现的。接着作者就为ConvNext网络单独使用了一个下采样层，就是通过一个Laryer Normalization加上一个卷积核大小为2步距为2的卷积层构成。更改后准确率就提升到了82.0%。

?3.?ConvNeXt 的版本

对于ConvNeXt网络，作者提出了T/S/B/L四个版本，计算复杂度刚好和Swin Transformer中的T/S/B/L相似。

这四个版本的配置如下：

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)
其中C代表4个stage中输入的通道数，B代表每个stage重复堆叠block的次数。

4.ConvNeXt的整体结构

下图是另一个博主画的网络结构图，博文链接(6条消息) ConvNeXt网络详解_太阳花的小绿豆的博客-CSDN博客

ConvNeXt Block会发现其中还有一个Layer Scale操作（论文中并没有提到），其实它就是将输入的特征层乘上一个可训练的参数，该参数就是一个向量，元素个数与特征层channel相同，即对每个channel的数据进行缩放。

?5.?ConvNeXt网络模型代码

代码只能说太简单了，是个人都能看懂，就直接放代码了。

"""
original code from facebook research:
https://github.com/facebookresearch/ConvNeXt
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


def drop_path(x, drop_prob: float = 0., training: bool = False):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
    'survival rate' as the argument.

    """
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)


class LayerNorm(nn.Module):
    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
    with shape (batch_size, channels, height, width).
    官方实现的LN是默认对最后一个维度进行的，这里是对channel维度进行的，所以单另设一个类。
    """

    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape), requires_grad=True)
        self.bias = nn.Parameter(torch.zeros(normalized_shape), requires_grad=True)
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise ValueError(f"not support data format '{self.data_format}'")
        self.normalized_shape = (normalized_shape,)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.data_format == "channels_last":
            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
        elif self.data_format == "channels_first":
            # [batch_size, channels, height, width]
            mean = x.mean(1, keepdim=True)
            var = (x - mean).pow(2).mean(1, keepdim=True)
            x = (x - mean) / torch.sqrt(var + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x


class Block(nn.Module):
    r""" ConvNeXt Block. There are two equivalent implementations:
    (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
    We use (2) as we find it slightly faster in PyTorch

    Args:
        dim (int): Number of input channels.
        drop_rate (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_rate=0., layer_scale_init_value=1e-6):
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
        self.norm = LayerNorm(dim, eps=1e-6, data_format="channels_last")
        self.pwconv1 = nn.Linear(dim, 4 * dim)  # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * dim, dim)
        # layer scale
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim,)),
                                  requires_grad=True) if layer_scale_init_value > 0 else None
        self.drop_path = DropPath(drop_rate) if drop_rate > 0. else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shortcut = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)  # [N, C, H, W] -> [N, H, W, C]
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2)  # [N, H, W, C] -> [N, C, H, W]

        x = shortcut + self.drop_path(x)
        return x


class ConvNeXt(nn.Module):
    r""" ConvNeXt
        A PyTorch 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.
        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
    """
    def __init__(self, in_chans: int = 3, num_classes: int = 1000, depths: list = None,
                 dims: list = None, drop_path_rate: float = 0., layer_scale_init_value: float = 1e-6,
                 head_init_scale: float = 1.):
        super().__init__()
        self.downsample_layers = nn.ModuleList()  # stem and 3 intermediate downsampling conv layers
        # stem为最初的下采样部分
        stem = nn.Sequential(nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
                             LayerNorm(dims[0], eps=1e-6, data_format="channels_first"))
        self.downsample_layers.append(stem)

        # 对应stage2-stage4前的3个downsample
        for i in range(3):
            downsample_layer = nn.Sequential(LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
                                             nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2))
            self.downsample_layers.append(downsample_layer)

        self.stages = nn.ModuleList()  # 4 feature resolution stages, each consisting of multiple blocks
        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
        cur = 0
        # 构建每个stage中堆叠的block
        for i in range(4):
            stage = nn.Sequential(
                *[Block(dim=dims[i], drop_rate=dp_rates[cur + j], layer_scale_init_value=layer_scale_init_value)
                  for j in range(depths[i])]
            )
            self.stages.append(stage)
            cur += depths[i]

        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)  # final norm layer
        self.head = nn.Linear(dims[-1], num_classes)
        self.apply(self._init_weights)
        self.head.weight.data.mul_(head_init_scale)
        self.head.bias.data.mul_(head_init_scale)

    def _init_weights(self, m):
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            nn.init.trunc_normal_(m.weight, std=0.2)
            nn.init.constant_(m.bias, 0)

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        for i in range(4):
            x = self.downsample_layers[i](x)
            x = self.stages[i](x)

        return self.norm(x.mean([-2, -1]))  # global average pooling, (N, C, H, W) -> (N, C)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.forward_features(x)
        x = self.head(x)
        return x


def convnext_tiny(num_classes: int):
    # https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth
    model = ConvNeXt(depths=[3, 3, 9, 3],
                     dims=[96, 192, 384, 768],
                     num_classes=num_classes)
    return model


def convnext_small(num_classes: int):
    # https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth
    model = ConvNeXt(depths=[3, 3, 27, 3],
                     dims=[96, 192, 384, 768],
                     num_classes=num_classes)
    return model


def convnext_base(num_classes: int):
    # https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth
    # https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth
    model = ConvNeXt(depths=[3, 3, 27, 3],
                     dims=[128, 256, 512, 1024],
                     num_classes=num_classes)
    return model


def convnext_large(num_classes: int):
    # https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth
    # https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth
    model = ConvNeXt(depths=[3, 3, 27, 3],
                     dims=[192, 384, 768, 1536],
                     num_classes=num_classes)
    return model


def convnext_xlarge(num_classes: int):
    # https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth
    model = ConvNeXt(depths=[3, 3, 27, 3],
                     dims=[256, 512, 1024, 2048],
                     num_classes=num_classes)
    return model