何凯明最新一作MAE解读系列2之代码实践

目录

  • 前言
  • 一睹为快,眼见为实
  • 代码解读


前言

github:https://github.com/facebookresearch/mae
论文:https://arxiv.org/abs/2111.06377
解读:何凯明最新一作MAE解读系列1

MAE实践主要从几个方面展开:

  1. 复现代码;复现代码过程,去粗存精,保留软件环境说明,如pytorch+cuda版本等;从jupyter来可视化过程。
  2. 解读代码。解读代码过程,双管齐下,参照论文framework来解读代码。

“知行合一”,“纸上得来终觉浅,觉知此事要躬行”

一睹为快,眼见为实

在jupyter中先快速可视化MAE的工作效果,以下展示使用MAE重构遮挡75%比例的图像效果,不用GPU即可完成操作。

github
download code+model
Demo: Jupyter Visualize
Debug
Show

Bug: __init__() got an unexpected keyword argument 'qk_scale' #58

https://github.com/facebookresearch/mae/issues/58#

Jupyter通过命令行保存成python格式

Jupyter通过命令行保存成python格式,代码如下(示例):

try:   
    !jupyter nbconvert --to python mae_visualize.ipynb
    # python即转化为.py,script即转化为.html
    # file_name.ipynb即当前module的文件名
except:
    pass

Run MAE on the image

# make random mask reproducible (comment out to make it change)
torch.manual_seed(2)
print('MAE with pixel reconstruction:')
run_one_image(img, model_mae)

MAE with pixel reconstruction:

何凯明最新一作MAE解读系列2之代码实践_第1张图片
以上部分,可视化了MAE可视化的效果。主要由三个部分实现的,第一,随机采样:通过高比例如75%的遮挡像素块,留下可见的25%子块;第二,编码器将这1/4的子块进行表征学习,注意,没有任何的令牌输入到编码器中;第三,解码器将编码后的表征与3/4遮挡子块对应的令牌一起输入到解码器中,解码器的输出为归一化的像素值,他们将会重塑成有序的图像,即重构后的图片。
以上部分是预训练的过程。在图像推断或者识别中,去掉解码器,只保留编码器部分,并且输入到编码器的是一张完整的图片,不再遮挡掩码。
我们解读MAE的目的就是预训练后的模型,作为有效的骨架,来为下游任务提供服务。目前,原文中实现的下游任务有目标检测(Faster RCNN)和图像分割(UpperNet),我们的目标是将之应用到2D pose中。
所以我们的思路如下:首先爬到fine-tuned的分类模型,作为我们的backbone;第二,在backbone后设计一个light-weighted Head,输出heapmap以适应pose估计的任务。

代码解读

Mask Autoencoder结构代码如下:

def __init__(self, img_size=224, patch_size=16, in_chans=3,
                 embed_dim=1024, depth=24, num_heads=16,
                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
                 mlp_ratio=4., norm_layer=nn.LayerNorm, norm_pix_loss=False):
        super().__init__()

        # --------------------------------------------------------------------------
        # MAE encoder specifics
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)  # fixed sin-cos embedding

        self.blocks = nn.ModuleList([
            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
            for i in range(depth)])
        self.norm = norm_layer(embed_dim)
        # --------------------------------------------------------------------------

        # --------------------------------------------------------------------------
        # MAE decoder specifics
        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)

        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))

        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False)  # fixed sin-cos embedding

        self.decoder_blocks = nn.ModuleList([
            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, qk_scale=None, norm_layer=norm_layer)
            for i in range(decoder_depth)])

        self.decoder_norm = norm_layer(decoder_embed_dim)
        self.decoder_pred = nn.Linear(decoder_embed_dim, patch_size**2 * in_chans, bias=True) # decoder to patch
        # --------------------------------------------------------------------------

        self.norm_pix_loss = norm_pix_loss

        self.initialize_weights()

MAE结构,主要包含两个部分构成,即encoder+decoder定义。encoder结构,基本上就是ViT-Large模型,包含24个transformer block;decoder 是8个transformer block的轻量级网络。

有关PatchEmbed、Block和pos_embed的原理与用法,会专门设置一个主题来讲原始的transformer。

patchify函数:将N×3×224×224图像变成图像块,每个图像块大小为16,并将16×16图像拉伸为1×256序列,这样输出为:N×196×768

    def patchify(self, imgs):
        """
        imgs: (N, 3, H, W)
        x: (N, L, patch_size**2 *3)
        """
        p = self.patch_embed.patch_size[0]   #p=16
        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0

        h = w = imgs.shape[2] // p   # 224/16=14
        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
        x = torch.einsum('nchpwq->nhwpqc', x)
        x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3))   #14*14=196,16^2=256
        return x

unpatchify函数:patchify的逆过程,将序列恢复到图像。即序列N×196×768,恢复到RGB图像N×3×224×224。

    def unpatchify(self, x):
        """
        x: (N, L, patch_size**2 *3)
        imgs: (N, 3, H, W)
        """
        p = self.patch_embed.patch_size[0]   #16
        h = w = int(x.shape[1]**.5)  #14
        assert h * w == x.shape[1]
        
        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
        x = torch.einsum('nhwpqc->nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
        return imgs

random_masking函数:执行随机采样。随机采样的过程是,首先通过均匀分布产生一个随机的数组,然后将之进行高低排序,选择小数值即1/4部分保留,余下的去掉,即洗牌操作,可实现图像块的随机遮挡。

    def random_masking(self, x, mask_ratio):
        """
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [N, L, D], sequence
        """
        N, L, D = x.shape  # batch, length, dim=N×196×768
        len_keep = int(L * (1 - mask_ratio))  #保留的图像块个数为49
        
        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_keep = ids_shuffle[:, :len_keep]
        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))

        # generate the binary mask: 0 is keep, 1 is remove
        mask = torch.ones([N, L], device=x.device)
        mask[:, :len_keep] = 0
        # unshuffle to get the binary mask
        mask = torch.gather(mask, dim=1, index=ids_restore)

        return x_masked, mask, ids_restore

有关torch.gather,可以参考torch官方文档说明。dim=1即按照列向量组合index,去取值mask。

forward_encoder函数:一张完整的图片输入到encoder中,首先执行常见的patch_embedding,然后添加2D_sin_cos位置信息;再执行随机掩码,这里随机掩码后的输出x与原始输入x的维度是否一致,从代码中看,masked后的输出添加了class token和pos_embed,因此两者的维度应该一致。然后将class token与masked后的x合并起来,输入到一个堆叠的block栈中。最后需要进行Layer Norm归一化。

    def forward_encoder(self, x, mask_ratio):
        # embed patches
        x = self.patch_embed(x)

        # add pos embed w/o cls token
        x = x + self.pos_embed[:, 1:, :]

        # masking: length -> length * mask_ratio
        x, mask, ids_restore = self.random_masking(x, mask_ratio)

        # append cls token
        cls_token = self.cls_token + self.pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        # apply Transformer blocks
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)

        return x, mask, ids_restore

forward_decoder函数:输入到解码器中,包含两个部分,即encoded学习到的representation,和剩下3/4部分的图像块对应的掩码id。mask_token是个虚构的标记,文章中说是为了模仿NLP中的class_token而设计的一个虚构标记;然后将encoded的representation和mask_token通过tensor合并到一起,之后再添加class_token。这波操作,真的是看的人云里雾里。mask_token和class_token具体是怎样的,只能在debug中去打印出来,看下shape或者数据。然后再添加decoder的位置信息,输入到一个浅层的transformer栈中,最后除了layer norm外,还通过以全连接层映射到输出,并去掉class_token。不禁提问,为什么这么做?

    def forward_decoder(self, x, ids_restore):
        # embed tokens
        x = self.decoder_embed(x)

        # append mask tokens to sequence
        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)
        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)  # no cls token
        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle
        x = torch.cat([x[:, :1, :], x_], dim=1)  # append cls token

        # add pos embed
        x = x + self.decoder_pos_embed

        # apply Transformer blocks
        for blk in self.decoder_blocks:
            x = blk(x)
        x = self.decoder_norm(x)

        # predictor projection
        x = self.decoder_pred(x)

        # remove cls token
        x = x[:, 1:, :]

        return x

forward_loss函数:输出是重构的向量,原始的图片可以通过patchify将图像数据转换成重构向量一样的维度。这里,原始图片即目标,需要像素值的归一化,首先取均值再求标准差。损失函数即均方误差,并且只在3/4部分上求误差。

    def forward_loss(self, imgs, pred, mask):
        """
        imgs: [N, 3, H, W]
        pred: [N, L, p*p*3]
        mask: [N, L], 0 is keep, 1 is remove, 
        """
        target = self.patchify(imgs)
        if self.norm_pix_loss:
            mean = target.mean(dim=-1, keepdim=True)
            var = target.var(dim=-1, keepdim=True)
            target = (target - mean) / (var + 1.e-6)**.5

        loss = (pred - target) ** 2
        loss = loss.mean(dim=-1)  # [N, L], mean loss per patch

        loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
        return loss

你可能感兴趣的:(Computer,Vision,python,深度学习,计算机视觉)