YOLO11 改进 – 主干网络_ ConvNeXtV2全卷积掩码自编码器网络:轻量级纯卷积架构破解特征坍塌难题,提升特征多样性

前言

本文介绍了将ConvNeXt V2与YOLOv11相结合的方法。先前的ConvNeXt模型与自监督学习结合效果不佳,为此提出全卷积掩码自动编码器框架和全局响应归一化(GRN)层,形成ConvNeXt V2模型家族,显著提升了纯卷积神经网络在各类识别基准上的性能。我们将ConvNeXt V2引入YOLOv11,对相关代码进行了修改和注册,并配置了yolov11 - ConvNeXtV2.yaml文件。实验结果表明,该结合方式在目标检测任务中展现出一定效果。

文章目录: YOLOv11改进大全:卷积层、轻量化、注意力机制、损失函数、Backbone、SPPF、Neck、检测头全方位优化汇总

专栏链接: YOLOv11改进专栏

文章目录

[TOC]

介绍

摘要

受益于架构优化与表示学习框架的持续改进,视觉识别领域在2020年代初期实现了显著的现代化进程与性能跃升。以ConvNeXt为代表的现代卷积神经网络在各种应用场景中展现出卓越的性能表现。尽管此类模型最初专为基于ImageNet标签的监督学习范式设计,理论上具备从掩码自编码器等自监督学习技术中获益的潜力,但实证研究表明简单组合这两种方法的效果并不理想。为此,本文提出了一种完全卷积的掩码自编码器框架,并创新性地引入全局响应归一化层,该层可集成至ConvNeXt架构中以强化通道间特征竞争机制。这种自监督学习技术与架构改进的协同设计策略催生了新型模型家族ConvNeXt V2,该模型在多个识别基准测试中显著提升了纯卷积神经网络的性能上限,涵盖ImageNet图像分类、COCO目标检测以及ADE20K语义分割任务。研究团队同时提供了多尺度预训练ConvNeXt V2模型,包括参数量仅为3.7M、在ImageNet上达到76.7%顶级准确率的高效Atto模型,以及使用公开训练数据达到88.9%顶级准确率、参数量高达650M的Huge模型,为不同计算资源约束下的应用需求提供了完整解决方案。

文章链接

论文地址: 论文地址

代码地址: 代码地址

基本原理

先前的ConvNeXt 模型表现出很好的结果,但在与自我监督学习 (MAE) 相结合时表现不佳。ConvNeXt V2 通过整合全卷积掩码自动编码器框架和全局响应归一化 (GRN) 层来解决这个问题,从而提高了多个基准测试的性能。

全卷积掩码自动编码器

掩蔽

32×32 个色块中的 60% 将从原始图像中随机删除。仅使用随机调整大小的裁剪进行数据增强。

编码器设计

Conv Next 模型用作编码器。但是,在预训练期间,标准卷积层被转换为子流形解析卷积,以便模型仅对可见数据点进行操作。稀疏卷积可以在精车阶段转换回标准卷积,而无需任何额外处理。

解码器设计

一个轻量级的普通 ConvNeXt 块被用作解码器。

重建目标

均方误差是在重建图像的掩码块和块状归一化原始图像之间计算的。

实验设置

预训练和微调分别在ImageNet-1K数据集上进行了800和100个时期。

全局响应归一化

GRN 旨在提高通道的对比度和选择性。给定一个输入特征,GRN 单元执行三个步骤:1) 全局特征聚合,2) 特征归一化,以及 3) 特征校准。

在 ConvNeXt V2 中,GRN 层被添加到维度扩展 MLP 层之后。并且 LayerScale 在变得多余时被丢弃。

核心代码

class ConvNeXtV2(nn.Module):
    """ ConvNeXt V2

    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.
        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
    """
    def __init__(self, in_chans=3, num_classes=1000, 
                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], 
                 drop_path_rate=0., head_init_scale=1.
                 ):
        super().__init__()
        self.depths = depths
        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)
        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 residual blocks
        dp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] 
        cur = 0
        for i in range(4):
            stage = nn.Sequential(
                *[Block(dim=dims[i], drop_path=dp_rates[cur + j]) 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)):
            trunc_normal_(m.weight, std=.02)
            nn.init.constant_(m.bias, 0)

    def forward_features(self, x):
        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):
        x = self.forward_features(x)
        x = self.head(x)
        return x

实验

脚本

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
#     修改为自己的配置文件地址
    model = YOLO('/root/ultralytics-main/ultralytics/cfg/models/11/yolov11-ConvNeXtV2.yaml')
#     修改为自己的数据集地址
    model.train(data='/root/ultralytics-main/ultralytics/cfg/datasets/coco8.yaml',
                cache=False,
                imgsz=640,
                epochs=10,
                single_cls=False,  # 是否是单类别检测
                batch=8,
                close_mosaic=10,
                workers=0,
                optimizer='SGD',
                amp=True,
                project='runs/train',
                name='ConvNeXtV2',
                )

结果

THE END