ConvNeXt网络结构搭建

ConvNeXt-T 结构图

ConvNeXt Block模块搭建

class Block(nn.Module): # ConvNeXt Block模块
   
    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)  # 1x1的卷积层和全连接层的作用是一样的 pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU() # GELU激活函数
        self.pwconv2 = nn.Linear(4 * dim, dim) # 注意pwconv1和pwconv2的输入输出channel是不同的
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim,)), # layer_scale层
                                  requires_grad=True) if layer_scale_init_value > 0 else None
        self.drop_path = DropPath(drop_rate) if drop_rate > 0. else nn.Identity() # 构建DropPath层

    def forward(self, x: torch.Tensor) -> torch.Tensor: # 正向传播过程
        shortcut = x
        x = self.dwconv(x) # 通过DW卷积
        x = x.permute(0, 2, 3, 1)  # 通过permute方法调整通道顺序 [N, C, H, W] -> [N, H, W, C]
        x = self.norm(x) # LayerNorm层
        x = self.pwconv1(x) # 1x1的卷积层
        x = self.act(x) # GELU激活函数
        x = self.pwconv2(x) # 1x1的卷积层
        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) # 通过drop_path层并融合shortcut
        return x

ConvNeXt整体网络结构搭建

class ConvNeXt(nn.Module):
  
    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 = nn.Sequential(nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
                             LayerNorm(dims[0], eps=1e-6, data_format="channels_first")) # 构建卷积层和LayerNorm层
        self.downsample_layers.append(stem) # 添加到downsample_layers中

        # 对应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) # 添加到downsample_layers中

        self.stages = nn.ModuleList()  # 存储每一个stage所构建的block 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)  # 最后一个LayerNorm层 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)

    # 顺序为downsample_layers[0]->stages[0]->downsample_layers[1]->stages[1]->downsample_layers[2]->stages[2]->downsample_layers[3]->stages[3]
    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)

        # 再通过全局平均池化以及LayerNorm层
        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

对于ConvNeXt网络，作者提出了T/S/B/L/XL五个版本

其中C代表4个stage中输入的通道数，B代表每个stage重复堆叠block的次数

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

reference

ConvNeXt网络详解_太阳花的小绿豆的博客-CSDN博客

13.1 ConvNeXt网络讲解_哔哩哔哩_bilibili