代码地址:
GitHub - facebookresearch/ConvNeXt-V2: Code release for ConvNeXt V2 model
论文地址:
[2301.00808] ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders
一、简介
ConvNext V2是Facebook AI Research提出的ConvNeXT改进版本。在计算机视觉中,虽然基于Transformer的架构在多个任务中表现超过了纯卷积架构网络,但是作者通过借鉴Transformer的思想,充分吸收其关键设计,如分层特征图、大感受野、更少的激活函数等,最终提出了ConvNeXt网络,并且在多个任务上证明了纯卷积设计能够达到甚至超越Transformer的性能。
ConvNeXt V2在ConvNeXt的基础上,引入了自监督学习,其核心创新点在于其自监督学习技术和架构改进的协同设计(Co-desiging):
1、为ConvNet量身定制的全卷积自监督学习框架:全卷积掩码自编码器(FCMAE)
2、在架构上引入针对性改进,以解决自监督训练时出现的特定问题:提出GRN (Global Response Normalization) 层,用于增强通道间的特征竞争,模拟大脑神经元之间的横向抑制机制,提升特征的判别力。
全卷积掩码自编码器(FCMAE)
FCMAE通过随机遮蔽输入图像的部分区域,强迫模型从上下文信息中重建完整图像。使用子流形稀疏卷积代替标准卷积,使卷积操作仅在可见像素上发生,完美适配了卷积的密集滑动特性,避免了直接引入可学习掩码token带来的训练-测试不一致问题。
FCMAE 框架与传统的掩码自编码器(MAE)框架相比,FAMAE使用了全卷积结构,而不是使用全连接层来生成掩码和重建图像。稀疏卷积让卷积网络也能够高效、自然地处理稀疏输入,同时又天然保留图像的二维空间结构和局部性归纳偏置。同时使用多尺度掩码策略,而不是固定大小的掩码,可以增加模型对不同尺度特征的感知能力。
全局响应归一化(GRN)层
在FCMAE预训练过程中,研究者发现 ConvNeXt 模型存在“特征崩溃”(Feature Collapse)现象:许多特征图变得停滞或饱和,通道间的激活冗余,模型学习到的特征多样性不足。因此,他们提出了GRN (Global Response Normalization) 层,其目的是增强通道间的特征竞争,模拟大脑神经元之间的横向抑制机制,提升特征的判别力。
将GRN插入到ConvNeXt块的11点卷积之后,作用在扩展后的高维特征上,让模型能够更有效地利用通道维度,从而提取更丰富、更有区分度的特征。
GRN的实现代码如下所示,包含三个步骤:
- 空间聚合:对输入特征在空间维度(高 H 和宽 W)上计算L1范数,从而计算每个通道在整张图上的全局响应
- 通道归一化:对上一步得到的通道能量取平均值,然后用各通道的能量除以这个平均值,得到一个归一化的比率,消除了图像整体亮度和对比度的影响
- 特征校准:利用这个比率对原始特征进行调整,实现“强者更强,弱者更弱”的竞争机制。
#models/utils.py class GRN(nn.Module): """ GRN (Global Response Normalization) layer """ def __init__(self, dim): super().__init__() self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) def forward(self, x): Gx = torch.norm(x, p=2, dim=(1,2), keepdim=True) Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) return self.gamma * (x * Nx) + self.beta + x二、在自定义数据集上训练
从仓库下载完代码后,可以看到在README.md有关于不同类型的ConvNeXt V2模型的预训练权重。训练自己数据集时,使用预训练权重能够加速模型收敛,并得到更好的性能表现。同时,还包含了单卡和多卡的测试命令,在TRAINING.md里有单卡和多卡的预训练和微调的命令。
需要注意的是,仓库代码里默认的数据集是ImageNet数据集,而在作者给的训练命令里没有指明这个参数值,所以在使用了自己的数据集路径后,会出现最后输出的类别数目错误相关的错误。因为在datast.py里面作者写死了分类的类别数目。所以,最好将data_set改为"image_folder"。
#main_finetune.py parser.add_argument('--data_set', default='IMNET', choices=['CIFAR', 'IMNET', 'image_folder'], type=str, help='ImageNet dataset path') #dataset.py if args.data_set == 'CIFAR': dataset = datasets.CIFAR100(args.data_path, train=is_train, transform=transform, download=True) nb_classes = 100 elif args.data_set == 'IMNET': print("reading from datapath", args.data_path) root = os.path.join(args.data_path, 'train' if is_train else 'val') dataset = datasets.ImageFolder(root, transform=transform) nb_classes = 1000 elif args.data_set == "image_folder": root = args.data_path if is_train else args.eval_data_path dataset = datasets.ImageFolder(root, transform=transform) nb_classes = args.nb_classes assert len(dataset.class_to_idx) == nb_classes else: raise NotImplementedError() print("Number of the class = %d" % nb_classes)下面是在单卡上给出的微调和测试命令示例,
#微调 python main_finetune.py \ --model convnextv2_base \ --batch_size 16 --update_freq 4 \ --blr 1e-3 \ --epochs 500 \ --warmup_epochs 20 \ --layer_decay_type 'group' \ --layer_decay 0.6 \ --weight_decay 0.05 \ --drop_path 0.1 \ --reprob 0.25 \ --smoothing 0. \ --model_ema True --model_ema_eval True \ --use_amp True \ --finetune convnextv2_base_1k_224_ema.pt \ --data_path /your_path \ --output_dir save_results \ --nb_classes 2 --input_size 60 \ --data_set image_folder #测试 python main_finetune.py \ --model convnextv2_base \ --eval true \ --resume save_results/checkpoint-best.pth \ --input_size 60 \ --data_path /your_path \ --nb_classes 2 --data_set image_folder