Masked Autoencoders Are Scalable Vision Learners

Posted 2021-12-18 Daft shiner

作为学术菜鸡的我跪着看完了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的patch
Rearrange('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 = 32
torch.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