EasyCV带你复现更好更快的自监督算法-FastConvMAE

我报名参加金石计划1期挑战——切割10万奖池，这是我的第7篇文章，点击检查活动概况

作者： 夕陌、谦言、莫申童、临在

导读

自监督学习（Self-Supervised Learning）利用很多无标注的数据进行表征学习，在特定下流使命上对参数进行微调，极大降低了图画使命繁重的标注作业，节约很多人力成本。近年来，自监督学习在视觉范畴大放异彩，受到了越来越多的重视。在CV范畴出现了如SIMCLR、MOCO、SwAV、DINO、MoBY、MAE等一系列作业。其间MAE的表现尤为惊艳，大家都被MAE简洁高效的功能所吸引，纷繁在 MAE上进行改善，例如MixMIM，VideoMAE等作业。MAE详解请参阅往期文章：MAE自监督算法介绍和根据EasyCV的复现 。

ConvMAE是由上海人工智能试验室和mmlab联合宣布在NeurIPS2022的一项作业，与MAE比较，练习相同的epoch数， ImageNet-1K 数据集的finetune准确率提高了 1.4%，COCO2017数据集上微调 25 个 epoch比较微调100 个 epoch 的 MAE AP box提高2.9， AP mask提高2.2， 语义切割使命上比较MAE mIOU提高3.6%。在此根底上，作者提出了FastConvMAE，进一步优化了练习功能，仅预练习50个epoch，ImageNet Finetuning的精度就超越MAE预练习1600个epoch的精度0.77个点（83.6/84.37)。在检测使命上，精度也超越ViTDet和Swin。

EasyCV是阿里巴巴开源的根据Pytorch，以自监督学习和Transformer技能为核心的 all-in-one 视觉算法建模工具，掩盖干流的视觉建模使命例如图画分类，衡量学习，方针检测，实例/语音/全景切割、关键点检测等范畴，具有较强的易用性和扩展性，一起注重功能调优，旨在为社区带来更多更快更强的算法。

近期FastConvMAE作业在EasyCV框架内首次对外开源，本文将要点介绍ConvMAE和FastConvMAE的首要作业，以及对应的代码实现，最终提供具体的教程示例怎么进行FastConvMAE的预练习和下流使命的finetune。

ConvMAE

ConvMAE是由上海人工智能试验室和mmlab联合宣布在NeurIPS2022里的一项作业，ConvMAE的提出证明了运用部分概括偏置和多标准的金字塔结构，经过MAE的练习方法能够学习到更好的特征表示。该作业提出：

1. 运用block-wise mask战略来保证核算功率。
2. 输出编码器的多标准特征，一起捕获细粒度和粗粒度图画信息。

原文参阅：arxiv.org/abs/2205.03…

试验成果显示，上述两项战略是简洁而有用的，使得ConvMAE在多个视觉使命中比较MAE获得了显着提高。以ConvMAE-Base和MAE-Base比较为例：在图画分类使命上， ImageNet-1K 数据集的微调准确率提高了 1.4%；在方针检测使命上，COCO2017微调 25 个 epoch 的AP box到达53.2%，AP mask到达47.1%，与微调100 个 epoch 的 MAE-Base比较别离提高2.9% 和 2.2% ；在语义切割使命上，运用UperNet网络头，ConvMAE-Base在ADE20K上的mIoU到达51.7%，比较MAE-Base提高3.6%。

ConvMAE的总体流程

与MAE不同的是，ConvMAE的编码器将输入图画逐渐抽象为多标准token embedding，而解码器则重建被mask掉的tokens对应的像素。关于前面stage部分的高分辨率token embedding，选用卷积块对部分进行编码，关于后边的低分辨率token embedding，则运用transformer来聚合大局信息。因此，ConvMAE的编码器在不同阶段能够一起获得部分和大局信息，并生成多标准特征。

当前的masked auto encoding框架，如BEiT，SimMIM，所选用的mask战略不能直接用于ConvMAE，由于在后边的transformer阶段，一切的tokens都需要保存。这导致对大模型进行预练习的核算成本过高，失去了MAE在transformer编码器中省去masked tokens的功率优势。此外，直接运用convolution-transformer结构的编码器进行预练习会导致卷积部分由于随机的mask而造成预练习的信息泄露，因此也会降低预练习所得模型的质量。

针对这些问题，ConvMAE提出了混合convolution-transformer架构。ConvMAE选用分块mask战略 (block-wise masking strategy)：，首要随机在后期的获取transformer token中生成后期的mask，然后对mask固定方位逐渐进行上采样到早期卷积阶段的高分辨率。这样，后期处理的token能够彻底分离为masked tokens和visible tokens，然后并承继了MAE运用稀少encoder的核算功率。

下面将别离针对encoder、mask战略以及decoder部分打开介绍。

Encoder

如总体流程图所示，encoder包括3 个阶段，每个阶段输出的特征维度别离是：H/4 W/4, H/8 W/8, H/16 W/16，其间H W为输入图画分辨率。前两个是卷积阶段，运用卷积模块将输入转换为token embeddings E1 ∈ R^(H/4 W/4 C1) and E2 ∈ R^(H/8 W/8 C2) 。其间卷积模块用5 5的卷积代替self-attention操作。前两个阶段的感触野较小首要捕捉图画的部分特征，第三个阶段运用transformer模块，将粗粒度特征交融, 并将感触野扩展到整个图画，获得token embeddings E3 ∈ R(H/16 W/16 C3)。在每个阶段之间，运用stride为2的卷积对tokens进行下采样。

其他包括transformer的结构，如CPT、Container、Uniformer、CMT、Swin等，在第一阶段的输入用相对方位编码或零填充卷积代替肯定方位编码，而作者发现在第3个transformer stage中运用肯定方位编码可获得最优功能。class token也从编码器中移除。

Mask战略

MAE、BEiT等，对输入patch选用随机mask。但相同的战略不能直接应用于ConvMAE编码器：如果独立地从stage-1的H/4 W/4个tokens中随机抽取mask，将导致降采样后的stage-3的简直一切token都有部分可见信息，使得编码器不再稀少。因此作者提出，从stage-3的输入tokens中以相同份额 （例如75%）生成mask，再对mask上采样2倍和4倍，别离作为stage-2和stage-1的mask。这样，ConvMAE在3个阶段都只含有很少的（例如25%）可见token，然后使得预练习时编码器的功率不受影响。而解码器的使命e则保持相同，即重建编码进程中被mask掉的tokens。

一起，前2个阶段的5X5卷积操作会在masked patches的边缘处泄漏不可见token的重建答案。为了避免这种状况保证预练习的质量，作者在前两个阶段选用了masked convolution, 使被mask掉的区域不参加编码进程。

Decoder

原始MAE的decoder的输入以编码器的输出和mask掉的tokens作为输入，然后经过堆叠的transformer blocks进行图画重建。ConvMAE编码器获得多标准特征E1、E2、E3，一起捕获细粒度和粗粒度图画信息。为了更好地的预练习，作者经过stride-4和stride-2卷积将E1和E2下采样到E3的相同大小，并进行多标准特征交融，再经过一个linear层得到最终要输入给 decoder 的可见token。方针函数和MAE相同，仅选用MSE作为损失函数，核算预测向量和被mask掉像素值之前的MSE loss，即只考虑mask掉的patches的重建。

下流使命

预练习之后，ConvMAE能够输出多标准的特征用于检测切割使命。

检测使命中，先将第stage-3的输出特征E3经过2×2最大池化获得E4。由于ConvMAE stage-3有11个self-attention层(ConvMAE-base)，核算成本过高，作者参阅ViT的benchmark将stage-3中除第1、4、7、11之外的一切global self-attention layers替换为了Window size77 的 local self-attention 层。修正后的local self-attention仍然由预练习的global self-attention进行初始化。global transformer blocks之间共享global relative position bias，local transformer blocks之间共享local relative position bias，这样就大大减轻了stage-3的核算和GPU内存开销。然后将多标准特征E1、E2、E3、E4送入MaskRCNN head进行方针检测。

而切割使命保存了stage-3的结构。

Benchmark

图画分类

ConvMAE根据ImageNet-1K，mask掉25%的input token做预练习，Decoder部分是一个8层的transformer，embedding 维度是512，head是12个。预练习参数和分类finetuning成果如下：

BEiT预练习300个epoch，finetune的精度到达83.0%，linear-prob的精度是37.6%。与BEiT比较，ConVMAE仅需要25%的token和一个轻量级的decoder finetune可到达85%，linear-prob能够到达70.9%。与本来的MAE比较，预练习相同的1600个epoch，ConVMAE比MAE提高1.4个点。与SimMIM（backbone运用Swin-B）比较提高了1个点。

检测

作者用ConvMAE替换Mask-RCNN的backbone，加载ConvMAE的预练习模型练习COCO数据集。

与ViT在COCO数据集上finetune100个epoch的成果比较，ConVMAE仅finetune 25个epoch在APbox和APmask就提高了2.9和2.2个点。

与ViTDet和MIMDet比较，ConvMAE finetune epoch更少、参数更少，别离超越了它们2.0%和1.7%。

与Swin和MViTv2比较，在APbox/APmask，其功能别离高出4.0%/3.6%和2.2%/1.4%。

切割

作者用ConvMAE替换UperNet的backbone，加载ConvMAE的预练习模型练习ADE20K数据集。

从成果中能够看出，比较与DeiT, Swin，MoCo-v3等网络ConvMAE取得了更高的功能(51.7%)。标明ConvMAE的多标准特征大大缩小了预练习Backbone 和下流网络之间的传输差距。

Fast ConvMAE

ConvMAE虽然在分类、检测、切割等下流使命中有了精度提高，并解决了pretraining-finetuning 的差异问题，但是模型的预练习依然耗时，ConvMAE的成果中，模型预练习了1600个epoch，因此作者又在ConvMAE的根底之上做了进一步的功能优化，提出了Fast ConvMAE，FastConvMAE提出了mask互补和deocder交融的计划，来实现快速的mask建模计划，进一步缩短了预练习的时间，从本来预练习的1600epoch缩短到了50epoch。 FastConvMAE的正式论文作者会在未来发出。

首要，FastConvMAE创新地设计出decoder互相交融的Mixture of Reconstructor (MoR)，能够让masked patches从不同的tokenizer中学习到互补的信息，包括EMA 的self-ensembling性质，DINO的similarity-discrimination才能,以及CLIP的multimodal常识。MoR首要包括两个部分，Partially-Shared Decoder(PS-Decoder)和Mixture of Tokenizer(MoT), PS-Decoder能够避免不同tokenizer的不同常识之间会产生梯度的冲突，MoT是用来生成不同的token作为masked patches的target。

一起Mask部分选用了互补战略，本来的mask每次只会保存例如25%的tokens，FastConvMAE将mask分成了4份，每一份都保存25%，4份mask之间互补。这样，相当于1张图片被分成了4张图片进行学习，理论上到达了4倍的学习效果。

    def random_masking(self, x, mask_ratio=None):
        """
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [N, L, D], sequence
        """
        N = x.shape[0]
        L = self.num_patches
        len_keep = int(L * (1 - mask_ratio))
        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
        # sort noise for each sample
        ids_shuffle = torch.argsort(
            noise, dim=1)  # ascend: small is keep, large is remove
        ids_restore = torch.argsort(ids_shuffle, dim=1)
        # keep the first subset
        ids_keep1 = ids_shuffle[:, :len_keep]
        ids_keep2 = ids_shuffle[:, len_keep:2 * len_keep]
        ids_keep3 = ids_shuffle[:, 2 * len_keep:3 * len_keep]
        ids_keep4 = ids_shuffle[:, 3 * len_keep:]
        # generate the binary mask: 0 is keep, 1 is remove
        mask1 = torch.ones([N, L], device=x.device)
        mask1[:, :len_keep] = 0
        # unshuffle to get the binary mask
        mask1 = torch.gather(mask1, dim=1, index=ids_restore)
        mask2 = torch.ones([N, L], device=x.device)
        mask2[:, len_keep:2 * len_keep] = 0
        # unshuffle to get the binary mask
        mask2 = torch.gather(mask2, dim=1, index=ids_restore)
        mask3 = torch.ones([N, L], device=x.device)
        mask3[:, 2 * len_keep:3 * len_keep] = 0
        # unshuffle to get the binary mask
        mask3 = torch.gather(mask3, dim=1, index=ids_restore)
        mask4 = torch.ones([N, L], device=x.device)
        mask4[:, 3 * len_keep:4 * len_keep] = 0
        # unshuffle to get the binary mask
        mask4 = torch.gather(mask4, dim=1, index=ids_restore)
        return [ids_keep1, ids_keep2, ids_keep3,
                ids_keep4], [mask1, mask2, mask3, mask4], ids_restore

前两个卷积阶段将输入转换为embeddings tokens E1和E2。然后E1和E2别离从4份mask中获取4份可见的tokens并进行拼接，作为decoder的输入，Decoder处理的是拼接后的tokens。代码参阅如下：

   def encoder_forward(self, x, mask_ratio):
        # embed patches
        ids_keep, masks, ids_restore = self.random_masking(x, mask_ratio)
        mask_for_patch1 = [
          1 - mask.reshape(-1, 14, 14).unsqueeze(-1).repeat(
            1, 1, 1, 16).reshape(-1, 14, 14, 4, 4).permute(
            0, 1, 3, 2, 4).reshape(x.shape[0], 56, 56).unsqueeze(1)
          for mask in masks
        ]
        mask_for_patch2 = [
          1 - mask.reshape(-1, 14, 14).unsqueeze(-1).repeat(
            1, 1, 1, 4).reshape(-1, 14, 14, 2, 2).permute(
            0, 1, 3, 2, 4).reshape(x.shape[0], 28, 28).unsqueeze(1)
          for mask in masks
        ]
        s1 = self.patch_embed1(x)
        s1 = self.pos_drop(s1)
        for blk in self.blocks1:
            s1 = blk(s1, mask_for_patch1)
        s2 = self.patch_embed2(s1)
        for blk in self.blocks2:
            s2 = blk(s2, mask_for_patch2)
        stage1_embed = self.stage1_output_decode(s1).flatten(2).permute(0, 2, 1)
        stage2_embed = self.stage2_output_decode(s2).flatten(2).permute(0, 2, 1)
        stage1_embed_1 = torch.gather(
          stage1_embed,
          dim=1,
          index=ids_keep[0].unsqueeze(-1).repeat(1, 1, stage1_embed.shape[-1]))
        stage2_embed_1 = torch.gather(
          stage2_embed,
          dim=1,
          index=ids_keep[0].unsqueeze(-1).repeat(1, 1, stage2_embed.shape[-1]))
        stage1_embed_2 = torch.gather(
          stage1_embed,
          dim=1,
          index=ids_keep[1].unsqueeze(-1).repeat(1, 1, stage1_embed.shape[-1]))
        stage2_embed_2 = torch.gather(
          stage2_embed,
          dim=1,
          index=ids_keep[1].unsqueeze(-1).repeat(1, 1, stage2_embed.shape[-1]))
        stage1_embed_3 = torch.gather(
          stage1_embed,
          dim=1,
          index=ids_keep[2].unsqueeze(-1).repeat(1, 1, stage1_embed.shape[-1]))
        stage2_embed_3 = torch.gather(
          stage2_embed,
          dim=1,
          index=ids_keep[2].unsqueeze(-1).repeat(1, 1, stage2_embed.shape[-1]))
        stage1_embed_4 = torch.gather(
          stage1_embed,
          dim=1,
          index=ids_keep[3].unsqueeze(-1).repeat(1, 1, stage1_embed.shape[-1]))
        stage2_embed_4 = torch.gather(
          stage2_embed,
          dim=1,
          index=ids_keep[3].unsqueeze(-1).repeat(1, 1, stage2_embed.shape[-1]))
        stage1_embed = torch.cat([
          stage1_embed_1, stage1_embed_2, stage1_embed_3, stage1_embed_4
        ])
        stage2_embed = torch.cat([
          stage2_embed_1, stage2_embed_2, stage2_embed_3, stage2_embed_4
        ])
        x = self.patch_embed3(s2)
        x = x.flatten(2).permute(0, 2, 1)
        x = self.patch_embed4(x)
        # add pos embed w/o cls token
        x = x + self.pos_embed
        x1 = torch.gather(x, dim=1, index=ids_keep[0].unsqueeze(-1).repeat(1, 1, x.shape[-1]))
        x2 = torch.gather(x, dim=1, index=ids_keep[1].unsqueeze(-1).repeat(1, 1, x.shape[-1]))
        x3 = torch.gather(x, dim=1, index=ids_keep[2].unsqueeze(-1).repeat(1, 1, x.shape[-1]))
        x4 = torch.gather(x, dim=1, index=ids_keep[3].unsqueeze(-1).repeat(1, 1, x.shape[-1]))
        x = torch.cat([x1, x2, x3, x4])
        # apply Transformer blocks
        for blk in self.blocks3:
            x = blk(x)
        x = x + stage1_embed + stage2_embed
        x = self.norm(x)
        mask = torch.cat([masks[0], masks[1], masks[2], masks[3]])
        return x, mask, ids_restore

Benchmark

EasyCV复现的成果如下：

ImageNet Pretrained

Config	Epochs	Download
fast_convmae_vit_base_patch16_8xb64_50e	50	model- log

ImageNet Finetuning

Algorithm	Fintune Config	Pretrained Config	Top-1	Download
Fast ConvMAE（EasyCV)	fast_convmae_vit_base_patch16_8xb64_100e_fintune	fast_convmae_vit_base_patch16_8xb64_50e	84.4%	fintune model- log
Fast ConvMAE（官方）	84.4%

Object Detection

Algorithm	Eval Config	Pretrained Config	mAP (Box)	mAP (Mask)	Download
Fast ConvMAE（EasyCV)	mask_rcnn_conv_vitdet_50e_coco	fast_convmae_vit_base_patch16_8xb64_50e	51.3%	45.6%	finetune model
Fast ConvMAE（官方）	51.0%	45.4%

从成果能够看出，仅预练习50个epoch，ImageNet Finetuning的精度就超越MAE预练习1600个epoch的精度0.77个点（83.6/84.37)。在检测使命上，精度也超越ViTDet和Swin。

FastConvMAE的更多官方成果请参阅：github.com/Alpha-VL/Fa… 。

Tutorial

一、装置依靠包

如果是在本地开发环境运转，能够参阅该链接装置环境。若运用PAI-DSW进行试验则无需装置相关依靠，在PAI-DSW docker中已内置相关环境。

二、数据预备

数据预备请参阅文档：github.com/alibaba/Eas…

三、模型预练习

FastConvMAE占用显存较大，主张运用A100资源。（FastConvMAE一次forward-backward等价于ConvMAE forward-backward 4次）

在EasyCV中，运用装备文件的方法来实现对模型参数、数据输入及增广方法、练习战略的装备，仅经过修正装备文件中的参数设置，就能够完成试验装备进行练习。

装备EasyCV途径

# 检查easycv装置方位
import easycv
print(easycv.__file__)

$ export PYTHONPATH=$PYTHONPATH:${your EasyCV root path}

练习

$ python -m torch.distributed.launch --nproc_per_node=8 --master_port=29930 \
tools/train.py \
configs/selfsup/fast_convmae/fast_convmae_vit_base_patch16_8xb64_50e.py \
--work_dir ./work_dir \
--launcher pytorch

下流使命finetune

下载预练习模型

$ wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/selfsup/FastConvMAE/pretrained/epoch_50.pth

• 单卡

$ python tools/train.py \
${CONFIG_FILE} \
--work_dir ./work_dir \
--load_from=./epoch_50.pth

• 多卡

$ python -m torch.distributed.launch --nproc_per_node=8 --master_port=29930 \
tools/train.py \
${CONFIG_FILE} \
--work_dir ./work_dir \
--launcher pytorch \
--load_from=./epoch_50.pth

分类使命 CONFIG_FILE 请参阅：github.com/alibaba/Eas…

分类使命 CONFIG_FILE 请参阅：github.com/alibaba/Eas…

Reference

EasyCV：github.com/alibaba/Eas…

EasyCV往期分享

• YOLOX-PAI:加速YOLOX,比YOLOV6更快更强
• 根据EasyCV复现DETR和DAB-DETR，Object Query的正确打开方法
• 根据EasyCV复现ViTDet：单层特征超越FPN
• MAE自监督算法介绍和根据EasyCV的复现
• EasyCV开源｜开箱即用的视觉自监督+Transformer算法库