Masked AutoEncoders（MAE) Top-1准确率87.8%

一种可扩展的计算机视觉自监督学习方法：MAE（masked autoencoders），Paddle版。

AI Studio

1363人浏览 · 2021-12-04 08:10:59

AI Studio · 2021-12-04 08:10:59 发布

Masked AutoEncoders（MAE) Top-1准确率87.8%

masked autoencoders（MAE）是一种可扩展的计算机视觉自监督学习方法。

本文的MAE方法很简单：mask输入图像的随机patch，并重建丢失的像素。它基于两个核心设计的。

首先，作者开发了一种非对称的编码器-解码器体系结构，其中的编码器只对可见的patch子集（没有mask的token）进行操作，同时还有一个轻量级的解码器，该解码器从潜在表示和mask token中重建原始图像。其次，作者发现mask高比例的输入图像patch（例如75%）会变成一个不错且有意义的自监督任务。将这两种设计结合起来，能够高效地训练大型模型：本文的模型加快训练速度（3倍或更多）并提高精度。

本文的可扩展方法使得学习具有良好通用性的高容量模型：例如，在仅使用ImageNet-1K数据时，ViT-Huge模型的Top-1精确度为87.8%。此外，下游任务中的迁移性能优于有监督的预训练。

此文最大的贡献，可能是在NLP和CV两大领域之间架起了一座更简便的桥梁。

此前，大名鼎鼎的GPT和BERT已经将大型自然语言处理（NLP）模型的性能提升到了一个新的高度。

直观点讲，就是事先遮住一些文本片段，让AI模型通过自监督学习，通过海量语料库的预训练，逐步掌握上下文语境，把这些被遮住的片段，用尽可能合乎逻辑的方式填回去。

这和我们做「完形填空」的方式有些类似。经过海量数据的学习和训练，AI模型慢慢学会了自己生成自然文本。目前，随着GPT及其后续改进模型的不断进步，生成的自然文本几乎可以乱真。

现在，何恺明的这篇文章把NLP领域已被证明极其有效的方式，用在了计算机视觉（CV）领域，而且模型更简单。

一、主要思路

MAE论文从三个方面做了分析，这也是MAE方法的立意：

图像的主流模型是CNN，而NLP的主流模型是transformer，CNN和transformer的架构不同导致NLP的BERT很难直接迁移到CV。但是vision transformer的出现已经解决这个问题；
图像和文本的信息密度不同，文本是高语义的人工创造的符号，而图像是一种自然信号，两者采用masked autoencoding建模任务难度就不一样，从句子中预测丢失的词本身就是一种复杂的语言理解任务，但是图像存在很大的信息冗余，一个丢失的图像块很容易利用周边的图像区域进行恢复；
用于重建的decoder在图像和文本任务发挥的角色有区别，从句子中预测单词属于高语义任务，encoder和decoder的gap小，所以BERT的decoder部分微不足道（只需要一个MLP），而对图像重建像素属于低语义任务（相比图像分类），decoder需要发挥更大作用：将高语义的中间表征恢复成低语义的像素值。

二、MAE模型结构

基于这三个的分析，论文提出了一种用于图像领域（ViT模型）的更简单有效的无监督训练方法：MAE（masked autoencoder），随机mask掉部分patchs然后进行重建，其整体架构如下所示。MAE采用encoder-decoder结构（分析3，需要单独的decoder），但属于非对称结构，一方面decoder采用比encoder更轻量级设计，另外一方面encoder只处理一部分patchs（visible patchs，除了masked patchs之外的patchs），而encoder处理所有的patchs。一个很重要的点，MAE采用很高的masking ratio（比如75%甚至更高），这契合分析2，这样构建的学习任务大大降低了信息冗余，也使得encoder能学习到更高级的特征。由于encoder只处理visible patchs，所以很高的masking ratio可以大大降低计算量。

在预训练期间，大比例的随机的图像块子集（如 75%）被屏蔽掉。编码器用于可见patch的小子集。在编码器之后引入掩码标记，并且完整的编码块和掩码标记集由一个小型解码器处理，该解码器以像素为单位重建原始图像。

预训练后，解码器被丢弃，编码器应用于未损坏的图像以生成识别任务的表示。

MAE 是一种简单的自编码方法，可以在给定部分观察的情况下重建原始信号。由编码器将观察到的信号映射到潜在表示，再由解码器从潜在表示重建原始信号。

与经典的自动编码器不同，MAE采用非对称设计，允许编码器仅对部分观察信号（无掩码标记）进行操作，并采用轻量级解码器从潜在表示和掩码标记中重建完整信号。

掩膜

将图像划分为规则的非重叠patch。对patch的子集进行采样并屏蔽剩余patch。采样策略很简单：均匀分布，简单称为“随机抽样”。

MAE 编码器

编码器仅适用于可见的、未屏蔽的patch。编码器通过添加位置嵌入的线性投影嵌入patch，然后通过一系列 Transformer 块处理结果集。编码器只对整个集合的一小部分（如 25%）进行操作。

被屏蔽的patch会被移除；不使用掩码令牌。这样可以节约计算资源，使用一小部分计算和内存来训练非常大的编码器。

MAE解码器

解码器的输入是完整的令牌集。每个掩码标记代表一个共享的、学习过的向量，表示存在要预测的缺失patch。

解码器仅在预训练期间用于执行图像重建任务。因此，它的设计可以独立于编码器。实验中使用的解码器更加轻量级。通过这种非对称设计，显著减少了预训练时间。

图像目标的重建

MAE 通过预测每个掩码块的像素值来重建输入图像。

解码器输出中的每个元素都是一个表示补丁的像素值向量。解码器的最后一层是线性投影，其输出通道的数量等于补丁中像素值的数量。解码器的输出被重新整形以形成重建的图像。

MAE 预训练实施效率高，实现方式简单，而且不需要任何专门的稀疏操作。

随着输入图像被遮住的比例升高，MAE的性能迅速上升，在75%左右达到最佳效果。

三、实验

#1.实验的代码基于PaddleViT，数据基于ILSVRC2012 ImageNet验证集，训练集划分为4万张图片，验证集1万张图片，分类还是1000类。
#!git clone https://hub.fastgit.org/BR-IDL/PaddleViT.git

#2.安装依赖库
!pip install yacs
!pip install yaml
!pip install einops

#3.掩码的生成：根据输入图像尺寸和掩码率来生成
#!cat /home/aistudio/work/MAE/mask.py

#4.基于PaddleViT的VisionTransformer修改MAE模型
#!cat  /home/aistudio/work/MAE/mae.py

#5.解压Imagenet Val集,作为训练和验证集。原数据集解压后的txt文件中路径有误，用work目录下的文件覆盖
!tar xvf /home/aistudio/data/data89857/ILSVRC2012mini.tar -C /home/aistudio/data/ > /dev/null
!/bin/cp  –r –f /home/aistudio/work/train_list.txt /home/aistudio/data/ILSVRC2012mini
!/bin/cp  –r –f /home/aistudio/work/val_list.txt /home/aistudio/data/ILSVRC2012mini

#6.模型和训练配置
#!cat  /home/aistudio/PaddleViT/image_classification/MAE/configs/mae_base_patch16_224.yaml

#7.预训练程序：使用图像生成的标签，Loss使用MSELoss
#!cat /home/aistudio/PaddleViT/image_classification/MAE/main_single_gpu_pretrain.py

#8.启动预训练

%cd /home/aistudio/work/MAE
!python main_single_gpu_pretrain.py \
   -cfg='./configs/mae_base_patch16_224.yaml' \
   -dataset='imagenet2012' \
   -batch_size=128 \
   -data_path='/home/aistudio/data/ILSVRC2012mini'

# 9.预训练模型可视化，原图片、重建图片和掩码图片存在./vis目录下
'''
%cd /home/aistudio/work/MAE

!python vis.py -cfg='./configs/mae_base_patch16_224.yaml' \
-dataset='imagenet2012'\
-img_path='/home/aistudio/data/ILSVRC2012mini/val/ILSVRC2012_val_00040001.JPEG'\
-save_path='./vis'\
-model_path='/home/aistudio/work/MAE/output/train-20211124-18-14-37/MAE-Epoch-60-Loss-0.7674374569892883.pdparams' 
'''

'''
from PIL import Image
import matplotlib.pyplot as plt
%matplotlib inline    
ori_img =Image.open('vis/ori_img.jpg')
rec_img =Image.open('vis/rec_img.jpg')
mask_img =Image.open('vis/mask_img.jpg')
plt.figure(figsize=(15, 45))
plt.subplot(131);plt.imshow(ori_img);plt.xticks([]);plt.yticks([]);plt.title('Ori Image',color='red')
plt.subplot(132);plt.imshow(rec_img);plt.xticks([]);plt.yticks([]);plt.title('Rec Image',color='red')
plt.subplot(133);plt.imshow(mask_img);plt.xticks([]);plt.yticks([]);plt.title('Mask Image',color='red')
plt.show()
'''

在这里插入图片描述

#10.利用预训练得到的模型，进行finetune的配置。
#!cat  /home/aistudio/work/MAE/configs/mae_base_patch16_224_finetune.yaml

#11.finetune程序，标签采用数据文件中的标签，Loss采用CrossEntropyLoss，与一般ViT的FineTune方式一致。
#!cat /home/aistudio/work/MAE/main_single_gpu_finetune.py

#12.进行MAE FineTue,注意修改PRETRAINED(-pretrained)的模型所在位置
'''
%cd /home/aistudio/work/MAE
!python main_single_gpu_finetune.py \
   -cfg='./configs/mae_base_patch16_224_finetune.yaml' \
   -dataset='imagenet2012' \
   -batch_size=64 \
   -data_path='/home/aistudio/data/ILSVRC2012mini' \
   -pretrained='./output/train-20211124-18-14-37/MAE-Epoch-60-Loss-0.7674374569892883'
'''

# 13.基于ViT模型进行FineTune,注意修改用于评估模型（-pretrained）所在位置
'''
%cd /home/aistudio/work/MAE
!python main_single_gpu_finetune-VIT.py -cfg='./configs/vit_base_patch16_224.yaml' -dataset='imagenet2012' -batch_size=64 -data_path='/home/aistudio/data/ILSVRC2012mini' -pretrained='./output/train-20211124-18-14-37/MAE-Epoch-60-Loss-0.7674374569892883'
'''

^C

# 14.MAE模型评估,注意修改用于评估模型（-pretrained）所在位置
'''
%cd /home/aistudio/work/MAE
!python main_single_gpu_finetune.py \
-cfg='./configs/mae_base_patch16_224_finetune.yaml' \
-dataset='imagenet2012' \
-batch_size=64 \
-data_path='/home/aistudio/data/ILSVRC2012mini' \
-eval \
-pretrained='./output/train-20211123-17-27-45/MAE-Epoch-20-Loss-5.221875454711914'
'''