深度学习模型之CNN（二十五）ConvNeXt网络讲解及使用Pytorch搭建

网络结构学习

前言

自从ViT(Vision Transformer)在CV领域大放异彩，越来越多的研究人员开始拥入Transformer的怀抱，而卷积神经网络已经开始慢慢淡出舞台中央。20221月，Facebook AI Research和UC Berkeley一起发表了一篇文章A ConvNet for the 2020s，在文章中提出了ConvNeXt纯卷积神经网络，它对标的是2021年非常火的Swin Transformer，通过一系列实验比对，在相同的FLOPs下，ConvNeXt相比Swin Transformer拥有更快的推理速度以及更高的准确率，在ImageNet 22K上ConvNeXt-XL达到了87.8%的准确率（下图所示）。

ConvNeXt其实“毫无亮点”，因为它使用的全部都是现有的结构和方法，没有任何结构或者方法的创新。而且源码也非常的精简，100多行代码就能搭建完成。而Swin Transformer的滑动窗口，相对位置索引等不仅原理理解起来很吃力，源码也非常多，但Swin Transformer确实十分成功并且设计的非常巧妙。

作者认为，基于Transformer架构的模型效果比卷积神经网络要好的原因可能在于随着技术的不断发展，各种新的架构以及优化策略促使Transformer模型的效果更好。所以作者想要使用相同的策略去训练卷积神经网络看看能不能达到和Transformer一样的效果，因此做了一系列实验。

设计方案

作者首先利用训练vision Transformers的策略去训练原始的ResNet50模型，得出效果为78.8，并将此结果作为后续实验的基准baseline。如下图所示，ConvNeXt-T/B达到的精度有82.0%，Swin-T/B为81.3%。

ConvNeXt设计与实验：

• Macro design
• ResNeXt
• Inverted bottleneck
• Large kerner size
• Various layer-wise micro designs

Macro design

如上图所示，分为两个小部分，分别是stage ratio和“patchify” stem。对应的举措分别是讲ResNet-50中堆叠的底数由（3，4，6，3）调整为（3，3，9，3）（为了能和Swin Tranformer中的比例保持一致）、将stem换成卷积核大小为4且步距为4的卷积层。

Changing stage compute ratio，在原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中的堆叠次数由(3, 4, 6, 3)调整成(3, 3, 9, 3)，和Swin-T拥有相似的FLOPs。进行调整后，准确率由78.8%提升到了79.4%。

Changing stem to “Patchify”，在之前的卷积神经网络中，一般最初的下采样模块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也降低了一点。

ResNeXt

Depth conv

作者借鉴了ResNeXt中的组卷积grouped convolution，因为ResNeXt相比普通的ResNet而言在FLOPs以及accuracy之间做到了更好的平衡。这里作者采用的是更激进的depthwise convolution（groups数和输入特征矩阵的深度相等），即group数和通道数channel相同，且作者认为depthwise convolution和self-attention中的加权求和操作很相似。

下图（上左）所采用的是ResNet所采用的瓶颈结构，即两头粗中间细的结构（256->64->256）。而在ResNeXt中采用的是下图（上右）的结构。二者之间唯一的区别就是在中间3x3的卷积层部分，在ResNet中，3x3的卷积层就是一个普通卷积层，但在ResNeXt中采用的是组卷积。

DW卷积如下图（下）所示，这里就不细讲了，可以回到ResNeXt中去看看详细过程。

通过上述改成激进的DW卷积之后，准确率从79.5%降到了78%。

增大width（图片channel），将ResNet50对标Swin-T，前者第一个Stage输入特征矩阵的深度为64，后者第一个Stage输入的channel为96。因此，作者将ResNet的每一个Stage的输入特征矩阵的channel与Swin-T中每个Stage的输入特征矩阵的channel保持一致。于是准确率从78%升到了80.5%，虽然同时FLOPs也增长了很多。

Inverted Bottleneck

作者认为Transformer block中的MLP模块非常像MobileNetV2中的Inverted Bottleneck模块，即两头细中间粗。下图a是ReNet中采用的Bottleneck模块，b是MobileNetV2采用的Inverted Botleneck模块，c是ConvNeXt采用的是Inverted Bottleneck模块。

作者采用Inverted Bottleneck模块后，在较小的模型上准确率由80.5%提升到了80.6%，在较大的模型上准确率由81.9%提升到82.6%。

Large Kernel Sizes

在Transformer中一般都是对全局做self-attention，比如Vision Transformer。即使是Swin Transformer也有7x7大小的窗口。但现在主流的卷积神经网络都是采用3x3大小的窗口，因为之前VGG论文中说通过堆叠多个3x3的窗口可以替代一个更大的窗口，而且现在的GPU设备针对3x3大小的卷积核做了很多的优化，所以会更高效。接着作者做了如下两个改动：

Moving up depthwise conv layer，即将depthwise conv模块上移，原来是先通过1x1 conv -> 再depthwise conv ->然后 1x1 conv，现在变成了depthwise conv -> 1x1 conv -> 1x1 conv。这么做是因为在Transformer中，MSA模块是放在MLP模块之前的，所以这里进行效仿，将depthwise conv上移（也就是上图c）。这样改动后，准确率由80.6%下降到了79.9%，同时FLOPs也减小了。

Increasing the kernel size，接着将depthwise conv的卷积核大小由3x3改成了7x7（和Swin Transformer一样），作者也做了一系列实验，除了3x3之外，还包括 5, 7, 9, 11发现取到7时准确率就达到了饱和。并且准确率从79.9% (3×3) 增长到 80.6% (7×7)（正好和Swin Transformer里的窗口大小是一致的，简直是玄学）。

Micro Design

接下来作者在聚焦到一些更细小的差异，比如激活函数将ReLU更改为GELU以及Normalization。

• Replacing ReLU with GELU，在Transformer中激活函数基本用的都是GELU，而在卷积神经网络中最常用的是ReLU，于是作者将激活函数替换成了GELU，发现准确率没变化。
• Fewer activation functions，使用更少的激活函数。在卷积神经网络中，一般会在每个卷积层或全连接后都接上一个激活函数。但在Transformer中并不是每个模块后都跟有激活函数（如下图），比如MLP中只有第一个全连接层后跟了GELU激活函数。接着作者在ConvNeXt Block中也减少激活函数的使用，如下图所示，减少后发现准确率从80.6%增长到81.3%。
• Fewer normalization layers，使用更少的Normalization Layer。同样在Transformer中，Normalization使用的也比较少，接着作者也减少了ConvNeXt Block中的Normalization层，只保留了depthwise conv后的Normalization层。此时准确率已经达到了81.4%，已经超过了Swin-T。
• Substituting BN with LN，将BN替换成LN。Batch Normalization（BN）在卷积神经网络中是非常常用的模块，它可以加速网络的收敛并减少过拟合（但用的不好也是个大坑）。但在Transformer中基本都用的Layer Normalization（LN），因为最开始Transformer是应用在NLP领域的，BN又不适用于NLP相关任务。接着作者将BN全部替换成了LN，发现准确率提升达到了81.5%。

• Separate downsampling layers，单独的下采样层。在ResNet网络中stage2-4的第一个block都具有下采样功能（下图左），且都是通过将主分支上3x3的卷积层步距设置成2，捷径分支上1x1的卷积层步距设置成2进行下采样的。但在Swin Transformer中是通过一个单独的Patch Merging实现的（下图右）。接着作者就为ConvNext网络单独使用了一个下采样层，就是通过一个Laryer Normalization加上一个卷积核大小为2步距为2的卷积层构成。更改后准确率就提升到了82.0%，对应的就是ConvNeXt-T这个精度。

ConvNeXt参数

对于ConvNeXt网络，作者提出了T/S/B/L四个版本，计算复杂度刚好和Swin Transformer中的T/S/B/L相似。C代表4个stage中输入的通道数，B代表每个stage重复堆叠block的次数。

• 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-T结构图

up根据官方源码手绘的ConvNeXt-T网络结构图，仔细观察ConvNeXt Block会发现其中还有一个Layer Scale操作（论文中并没有提到），其实它就是将输入的特征层乘上一个可训练的参数，该参数就是一个向量，元素个数与特征层channel相同，即对每个channel的数据进行缩放。

工程目录

├── Test14_ConvNeXt
	├── model.py（模型文件）   
	├── my_dataset.py（数据处理文件）    
	├── train.py（调用模型训练，自动生成class_indices.json,ConvNeXt.pth文件）
	├── predict.py（调用模型进行预测）
	├── utils.py（工具文件，用得上就对了）  
	├── tulip.jpg（用来根据前期的训练结果来predict图片类型）
	└── convnext_tiny_1k_224_ema.pth（迁移学习，提前下载好convnext_tiny_1k_224_ema.pth权重脚本）
└── data_set
	└── data数据集

搭建网络

"""
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):
    # ......


class DropPath(nn.Module):
    def __init__(self, drop_prob=None):
        # ......

    def forward(self, x):
        # ......


class LayerNorm(nn.Module):
    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        # ......

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # ......


class Block(nn.Module):
    def __init__(self, dim, drop_rate=0., layer_scale_init_value=1e-6):
        # ......

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # ......


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.):
        # ......

    def _init_weights(self, m):
        # ......

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        # ......

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # ......


def convnext_tiny(num_classes: int):
    # ......


def convnext_small(num_classes: int):
    # ......


def convnext_base(num_classes: int):
    # ......


def convnext_large(num_classes: int):
    # ......


def convnext_xlarge(num_classes: int):
    # ......

模型文件

DropPath类

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)

LayerNorm类

实际上，官方有现成的LayerNorm方法，但是Pytorch实现的LayerNorm方法默认是从最后一个维度开始做Normalization。在ConvNeXt网络当中，是对channel维度进行LayerNorm处理的，如果说channel是放在最后一个维度的话，就可以用官方给的LayerNorm方法，但如果channel这个维度没有放在最后的话，就不能使用官方的方法，所以作者重写了LayerNorm方法。

data_format有两种形式，分别为channels_last和channels_first，分别对应着channel这个维度放在最后面（batch_size, height, width, channels）和没有放在最后一个位置（batch_size, channels, height, width，这通常是Pytorch默认的通道排列顺序）。

weight和bias分别对应着LayerNorm过程中的$\alpha$和$\beta$。

如果data_format是channels_first时，作者重写了这块的代码，首先对channel这个维度求均值（也就是对应dim为1的位置）得到mean，接下来求方差var，以及标准差处理。

最后再乘以weight并加上bias。

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).
    """

    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

Block类

self.gamma对应的是上图的Layer Scale，元素个数和输入特征层channel的个数是相同的（其实它就是将输入的特征层乘上一个可训练的参数，该参数就是一个向量，元素个数与特征层channel相同，即对每个channel的数据进行缩放）

在正向传播过程中，通过第一个permute方法调换Pytorch默认的通道排列顺序，即[N, C, H, W] -> [N, H, W, C]，再通过第二个permute方法将修改后的通道排列顺序给还原回去，即[N, H, W, C] -> [N, C, H, W]

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)
        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

ConvNeXt类

ConvNeXt-T: C = (96, 192, 384, 768), B = (3, 3, 9, 3)

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 = 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代表在当前Stage之前构建好了的block的个数
            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)
        # 对head，也就是Linear层对weight和bias乘上了head_init_scale（默认为1）
        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)

        # 相当于做了个Goble Avg Pooling操作
        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: 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)

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 

train

import os
import argparse

import torch
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms

from my_dataset import MyDataSet
from model import convnext_tiny as create_model
from utils import read_split_data, create_lr_scheduler, get_params_groups, train_one_epoch, evaluate


def main(args):
    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    print(f"using {device} device.")

    if os.path.exists("./weights") is False:
        os.makedirs("./weights")

    tb_writer = SummaryWriter()

    train_images_path, train_images_label, val_images_path, val_images_label = read_split_data(args.data_path)

    img_size = 224
    data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(img_size),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
        "val": transforms.Compose([transforms.Resize(int(img_size * 1.143)),
                                   transforms.CenterCrop(img_size),
                                   transforms.ToTensor(),
                                   transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}

    # 实例化训练数据集
    train_dataset = MyDataSet(images_path=train_images_path,
                              images_class=train_images_label,
                              transform=data_transform["train"])

    # 实例化验证数据集
    val_dataset = MyDataSet(images_path=val_images_path,
                            images_class=val_images_label,
                            transform=data_transform["val"])

    batch_size = args.batch_size
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
    print('Using {} dataloader workers every process'.format(nw))
    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batch_size,
                                               shuffle=True,
                                               pin_memory=True,
                                               num_workers=nw,
                                               collate_fn=train_dataset.collate_fn)

    val_loader = torch.utils.data.DataLoader(val_dataset,
                                             batch_size=batch_size,
                                             shuffle=False,
                                             pin_memory=True,
                                             num_workers=nw,
                                             collate_fn=val_dataset.collate_fn)

    model = create_model(num_classes=args.num_classes).to(device)

    if args.weights != "":
        assert os.path.exists(args.weights), "weights file: '{}' not exist.".format(args.weights)
        weights_dict = torch.load(args.weights, map_location=device)["model"]
        # 删除有关分类类别的权重
        for k in list(weights_dict.keys()):
            if "head" in k:
                del weights_dict[k]
        print(model.load_state_dict(weights_dict, strict=False))

    if args.freeze_layers:
        for name, para in model.named_parameters():
            # 除head外，其他权重全部冻结
            if "head" not in name:
                para.requires_grad_(False)
            else:
                print("training {}".format(name))

    # pg = [p for p in model.parameters() if p.requires_grad]
    pg = get_params_groups(model, weight_decay=args.wd)
    optimizer = optim.AdamW(pg, lr=args.lr, weight_decay=args.wd)
    lr_scheduler = create_lr_scheduler(optimizer, len(train_loader), args.epochs,
                                       warmup=True, warmup_epochs=1)

    best_acc = 0.
    for epoch in range(args.epochs):
        # train
        train_loss, train_acc = train_one_epoch(model=model,
                                                optimizer=optimizer,
                                                data_loader=train_loader,
                                                device=device,
                                                epoch=epoch,
                                                lr_scheduler=lr_scheduler)

        # validate
        val_loss, val_acc = evaluate(model=model,
                                     data_loader=val_loader,
                                     device=device,
                                     epoch=epoch)

        tags = ["train_loss", "train_acc", "val_loss", "val_acc", "learning_rate"]
        tb_writer.add_scalar(tags[0], train_loss, epoch)
        tb_writer.add_scalar(tags[1], train_acc, epoch)
        tb_writer.add_scalar(tags[2], val_loss, epoch)
        tb_writer.add_scalar(tags[3], val_acc, epoch)
        tb_writer.add_scalar(tags[4], optimizer.param_groups[0]["lr"], epoch)

        if best_acc < val_acc:
            torch.save(model.state_dict(), "./weights/best_model.pth")
            best_acc = val_acc


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_classes', type=int, default=5)
    parser.add_argument('--epochs', type=int, default=3)
    parser.add_argument('--batch-size', type=int, default=4)
    parser.add_argument('--lr', type=float, default=5e-4)
    parser.add_argument('--wd', type=float, default=5e-2)

    # 数据集所在根目录
    # https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz
    parser.add_argument('--data-path', type=str,
                        default="D:/python_test/deep-learning-for-image-processing/data_set/flower_data/flower_photos")

    # 预训练权重路径，如果不想载入就设置为空字符
    # 链接: https://pan.baidu.com/s/1aNqQW4n_RrUlWUBNlaJRHA  密码: i83t
    parser.add_argument('--weights', type=str, default='./convnext_tiny_1k_224_ema.pth',
                        help='initial weights path')
    # 是否冻结head以外所有权重
    parser.add_argument('--freeze-layers', type=bool, default=False)
    parser.add_argument('--device', default='cuda:0', help='device id (i.e. 0 or 0,1 or cpu)')

    opt = parser.parse_args()

    main(opt)

训练结果

predict

import os
import json

import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt

from model import convnext_tiny as create_model


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"using {device} device.")

    num_classes = 5
    img_size = 224
    data_transform = transforms.Compose(
        [transforms.Resize(int(img_size * 1.14)),
         transforms.CenterCrop(img_size),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

    # load image
    img_path = "tulip.jpg"
    assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
    img = Image.open(img_path)
    plt.imshow(img)
    # [N, C, H, W]
    img = data_transform(img)
    # expand batch dimension
    img = torch.unsqueeze(img, dim=0)

    # read class_indict
    json_path = './class_indices.json'
    assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)

    with open(json_path, "r") as f:
        class_indict = json.load(f)

    # create model
    model = create_model(num_classes=num_classes).to(device)
    # load model weights
    model_weight_path = "./weights/best_model.pth"
    model.load_state_dict(torch.load(model_weight_path, map_location=device))
    model.eval()
    with torch.no_grad():
        # predict class
        output = torch.squeeze(model(img.to(device))).cpu()
        predict = torch.softmax(output, dim=0)
        predict_cla = torch.argmax(predict).numpy()

    print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
                                                 predict[predict_cla].numpy())
    plt.title(print_res)
    for i in range(len(predict)):
        print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
                                                  predict[i].numpy()))
    plt.show()


if __name__ == '__main__':
    main()

预测结果