Masked Autoencoders Are Scalable Vision Learners

作为学术菜鸡的我跪着看完了kaiming大佬的论文，先po一个大佬主页：Kaiming He
在讲Masked Autoencoders Are Scalable Vision Learners这个之前，由于笔者对Transformer没有太深理解，因此会穿插一些transformer以及ViT的知识，那么接下来就废话不多说进入正题吧。
Masked Autoencoders Are Scalable Vision Learners

ViT

在讲MAE之前，为了更好的理解其思想，这里先简单的介绍下ViT。
ViT文章 和 ViT代码

ViT architecture

class ViT(nn.Module):    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = "cls", channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):        super().__init__()        image_height, image_width = pair(image_size)        patch_height, patch_width = pair(patch_size)        assert image_height % patch_height == 0 and image_width % patch_width == 0, "Image dimensions must be divisible by the patch size."        num_patches = (image_height // patch_height) * (image_width // patch_width)        patch_dim = channels * patch_height * patch_width        assert pool in {"cls", "mean"}, "pool type must be either cls (cls token) or mean (mean pooling)"        self.to_patch_embedding = nn.Sequential(            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),            nn.Linear(patch_dim, dim),  # dim = 1024        )        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))  # torch.Size([1, 65, 1024])        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))  # torch.Size([1, 1, 1024])        self.dropout = nn.Dropout(emb_dropout)        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)        self.pool = pool        self.to_latent = nn.Identity()        self.mlp_head = nn.Sequential(            nn.LayerNorm(dim),            nn.Linear(dim, num_classes)        )    def forward(self, img):        x = self.to_patch_embedding(img)        b, n, _ = x.shape  # torch.Size([1, 64, 1024])        cls_tokens = repeat(self.cls_token, "() n d -> b n d", b = b)        x = torch.cat((cls_tokens, x), dim=1)        x += self.pos_embedding[:, :(n + 1)]        x = self.dropout(x)        x = self.transformer(x)        x = x.mean(dim = 1) if self.pool == "mean" else x[:, 0]        x = self.to_latent(x)        return self.mlp_head(x)

先来看看einops如何实现patch的维度变化：PyTorch 70.einops：优雅地操作张量维度

# 3x256x256图片分为64个3x32x32的patchRearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width)"""torch.Size([1, 3, 256, 256])b = 1, c = 3, h = 8, p1 = 32, w = 8, p2 = 32torch.Size([1, 64, 3072])"""

经过Rearrange后就可以把原始图片分成多个patch，接着用nn.linear对其embedding成64x1024，接下来对其进行Positional Encoding（可学习的位置编码，为什么需要位置编码呢？详见Transformer Architecture: The Positional Encoding)和class_token（Vision Transformer）。然后送入transformer得到encoded embedding进行分类任务。

MAE architecture

有了ViT的前置知识后再来看MAE，其结构如下图所示：

其中ViT作为encoder，其输入的patches是没有经过mask的，注意这里虽然使用的那些patches是没有mask的，但是这些没有mask的只占所有patches的一小部分，这也就是为什么MAE能够只用较少的内存和计算消耗就能训练大的encoders。
接着decoder部分把所有patches都计算进去（包括编码后的patches和mask的patches），并加入了位置编码信息。这些mask的patches即要还原出来的图像。并使用mean squared error (MSE) 来计算重构图片和真实图片间的误差。

代码

此处代码是来自github别人的复现：Unofficial PyTorch implementation of Masked Autoencoders Are Scalable Vision Learners

def train_one_epoch(model: torch.nn.Module, data_loader: Iterable, optimizer: torch.optim.Optimizer,                    device: torch.device, epoch: int, loss_scaler, max_norm: float = 0, patch_size: int = 16,                     normlize_target: bool = True, log_writer=None, lr_scheduler=None, start_steps=None,                    lr_schedule_values=None, wd_schedule_values=None):    model.train()    metric_logger = utils.MetricLogger(delimiter="  ")    metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))    metric_logger.add_meter("min_lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))    header = "Epoch: [{}]".format(epoch)    print_freq = 10    loss_func = nn.MSELoss()    for step, (batch, _) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):        # assign learning rate & weight decay for each step        it = start_steps + step  # global training iteration        if lr_schedule_values is not None or wd_schedule_values is not None:            for i, param_group in enumerate(optimizer.param_groups):                if lr_schedule_values is not None:                    param_group["lr"] = lr_schedule_values[it] * param_group["lr_scale"]                if wd_schedule_values is not None and param_group["weight_decay"] > 0:                    param_group["weight_decay"] = wd_schedule_values[it]        images, bool_masked_pos = batch        images = images.to(device, non_blocking=True)        bool_masked_pos = bool_masked_pos.to(device, non_blocking=True).flatten(1).to(torch.bool)        # import pdb; pdb.set_trace()        with torch.no_grad():            # calculate the predict label            mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device)[None, :, None, None]            std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device)[None, :, None, None]            unnorm_images = images * std + mean  # in [0, 1]            if normlize_target:                images_squeeze = rearrange(unnorm_images, "b c (h p1) (w p2) -> b (h w) (p1 p2) c", p1=patch_size, p2=patch_size)                images_norm = (images_squeeze - images_squeeze.mean(dim=-2, keepdim=True)                    ) / (images_squeeze.var(dim=-2, unbiased=True, keepdim=True).sqrt() + 1e-6)                # we find that the mean is about 0.48 and standard deviation is about 0.08.                images_patch = rearrange(images_norm, "b n p c -> b n (p c)")            else:                images_patch = rearrange(unnorm_images, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=patch_size, p2=patch_size)            B, _, C = images_patch.shape            labels = images_patch[bool_masked_pos].reshape(B, -1, C)        with torch.cuda.amp.autocast():            outputs = model(images, bool_masked_pos)            loss = loss_func(input=outputs, target=labels)        loss_value = loss.item()        if not math.isfinite(loss_value):            print("Loss is {}, stopping training".format(loss_value))            sys.exit(1)        optimizer.zero_grad()        # this attribute is added by timm on one optimizer (adahessian)        is_second_order = hasattr(optimizer, "is_second_order") and optimizer.is_second_order        grad_norm = loss_scaler(loss, optimizer, clip_grad=max_norm,                                parameters=model.parameters(), create_graph=is_second_order)        loss_scale_value = loss_scaler.state_dict()["scale"]        torch.cuda.synchronize()        metric_logger.update(loss=loss_value)        metric_logger.update(loss_scale=loss_scale_value)        min_lr = 10.        max_lr = 0.        for group in optimizer.param_groups:            min_lr = min(min_lr, group["lr"])            max_lr = max(max_lr, group["lr"])        metric_logger.update(lr=max_lr)        metric_logger.update(min_lr=min_lr)        weight_decay_value = None        for group in optimizer.param_groups:            if group["weight_decay"] > 0:                weight_decay_value = group["weight_decay"]        metric_logger.update(weight_decay=weight_decay_value)        metric_logger.update(grad_norm=grad_norm)        if log_writer is not None:            log_writer.update(loss=loss_value, head="loss")            log_writer.update(loss_scale=loss_scale_value, head="opt")            log_writer.update(lr=max_lr, head="opt")            log_writer.update(min_lr=min_lr, head="opt")            log_writer.update(weight_decay=weight_decay_value, head="opt")            log_writer.update(grad_norm=grad_norm, head="opt")            log_writer.set_step()        if lr_scheduler is not None:            lr_scheduler.step_update(start_steps + step)    # gather the stats from all processes    metric_logger.synchronize_between_processes()    print("Averaged stats:", metric_logger)    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}

实验效果如下图所示，可以发现mask掉大部分的图片经过decoder后能还原出原始图像，但是随着mask rate的提高，其重构的图像还是能还原出学到的东西，只不过数量变少了。这些都是符合语义信息的（蘑菇还是蘑菇），说明模型已经学习到了图像中的物体归纳性特征，已经具有很强的泛化能力。

参考文献和资料：

1.Masked Autoencoders Are Scalable Vision Learners
2.ViT文章
3.ViT代码
4.PyTorch 70.einops：优雅地操作张量维度
5.Transformer Architecture: The Positional Encoding
6.Vision Transformer
7.Unofficial PyTorch implementation of Masked Autoencoders Are Scalable Vision Learners