带掩码的自编码器MAE详解和Pytorch代码实现

监督学习是训练机器学习模型的传统方法，它在训练时每一个观察到的数据都需要有标注好的标签。如果我们有一种训练机器学习模型的方法不需要收集标签，会怎么样?如果我们从收集的相同数据中提取标签呢?这种类型的学习算法被称为自监督学习。这种方法在自然语言处理中工作得很好。一个例子是BERT¹，谷歌自2019年以来一直在其搜索引擎中使用BERT¹。不幸的是，对于计算机视觉来说，情况并非如此。

Facebook AI的kaiming大神等人提出了一种带掩码自编码器(MAE)²，它基于(ViT)³架构。他们的方法在ImageNet上的表现要好于从零开始训练的VIT。在本文中，我们将深入研究他们的方法，并了解如何在代码中实现它。

带掩码自编码器(MAE)

对输入图像的patches进行随机掩码，然后重建缺失的像素。MAE基于两个核心设计。首先，开发了一个非对称的编码器-解码器架构，其中编码器仅对可见的patches子集(没有掩码的tokens)进行操作，同时还有一个轻量级的解码器，可以从潜在表示和掩码tokens重建原始图像。其次，发现对输入图像进行高比例的掩码，例如75%，会产生有意义的自监督任务。将这两种设计结合起来，能够高效地训练大型模型：加快模型训练速度(3倍甚至更多)并提高精度。

 

 

此阶段称为预训练，因为 MAE 模型稍后将用于下游任务，例如图像分类。模型在pretext上的表现在自监督中并不重要， 这些任务的重点是让模型学习一个预期包含良好语义的中间表示。在预训练阶段之后，解码器将被多层感知器 (MLP) 头或线性层取代，作为分类器输出对下游任务的预测。

 

 

 

模型架构

编码器

编码器是 ViT。它接受张量形状为 (batch_size, RGB_channels, height, width) 的图像。通过执行线性投影为每个Patch获得嵌入， 这是通过 2D 卷积层来完成。然后张量在最后一个维度被展平（压扁），变成 (batch_size, encoder_embed_dim, num_visible_patches)，并 转置为形状（batch_size、num_visible_patches、encoder_embed_dim）的张量。

 

完整代码请见原文：带掩码的自编码器MAE详解和Pytorch代码实现