ConvNeXt V2训练自定义分类数据集

代码地址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-desiging1、为ConvNet量身定制的全卷积自监督学习框架全卷积掩码自编码器FCMAE2、在架构上引入针对性改进以解决自监督训练时出现的特定问题提出GRN (Global Response Normalization) 层用于增强通道间的特征竞争模拟大脑神经元之间的横向抑制机制提升特征的判别力。全卷积掩码自编码器FCMAEFCMAE通过随机遮蔽输入图像的部分区域强迫模型从上下文信息中重建完整图像。使用子流形稀疏卷积代替标准卷积使卷积操作仅在可见像素上发生完美适配了卷积的密集滑动特性避免了直接引入可学习掩码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, p2, dim(1,2), keepdimTrue) Nx Gx / (Gx.mean(dim-1, keepdimTrue) 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, defaultIMNET, choices[CIFAR, IMNET, image_folder], typestr, helpImageNet dataset path) #dataset.py if args.data_set CIFAR: dataset datasets.CIFAR100(args.data_path, trainis_train, transformtransform, downloadTrue) 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, transformtransform) 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, transformtransform) 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