一文读懂MAE核心架构：从代码视角拆解MaskedAutoencoderViT工作机制

2026-02-05 04:18:10作者：庞眉杨Will

你是否在学习MAE（Masked Autoencoder for Visual Representation Learning）时，面对models_mae.py中复杂的类结构感到无从下手？本文将带你逐行解析这个视觉Transformer的核心实现，掌握掩码自编码器的工作原理。读完本文后，你将能够：

理解MAE的编码器-解码器架构设计
掌握图像分块与掩码处理的关键步骤
明晰前向传播中的数据流向
学会使用预定义模型配置快速上手

核心类结构概览

MAE的核心实现集中在MaskedAutoencoderViT类（定义于models_mae.py第22行），该类继承自PyTorch的nn.Module，采用典型的编码器-解码器架构。以下是其主要组件的关系示意图：

graph TD
    A[输入图像] --> B[PatchEmbed 分块嵌入]
    B --> C[位置嵌入]
    C --> D[随机掩码处理]
    D --> E[Transformer编码器块]
    E --> F[编码器输出]
    F --> G[解码器嵌入]
    G --> H[掩码标记填充]
    H --> I[解码器位置嵌入]
    I --> J[Transformer解码器块]
    J --> K[图像补丁预测]
    K --> L[损失计算]

编码器详解

图像分块与嵌入

编码器的首要任务是将输入图像转换为序列数据。在__init__方法（第33行）中，通过PatchEmbed类实现：

self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)

默认配置下，这会将224x224的RGB图像分割为16x16的补丁（patch），每个补丁通过线性投影转换为768维的嵌入向量（对应基础模型配置）。

位置嵌入与掩码机制

位置嵌入（第37行）采用固定的正弦余弦编码，通过util/pos_embed.py中的get_2d_sincos_pos_embed函数生成。掩码机制是MAE的核心创新点，由random_masking方法（第123行）实现：

def random_masking(self, x, mask_ratio):
    N, L, D = x.shape  # batch, length, dim
    len_keep = int(L * (1 - mask_ratio))
    
    noise = torch.rand(N, L, device=x.device)  # 生成随机噪声
    ids_shuffle = torch.argsort(noise, dim=1)  # 通过排序噪声确定掩码位置
    ids_keep = ids_shuffle[:, :len_keep]  # 保留的补丁索引
    x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
    
    # 生成掩码矩阵
    mask = torch.ones([N, L], device=x.device)
    mask[:, :len_keep] = 0
    mask = torch.gather(mask, dim=1, index=torch.argsort(ids_shuffle, dim=1))
    
    return x_masked, mask, ids_restore

默认情况下，75%的补丁会被掩码（通过mask_ratio=0.75控制），仅保留25%的补丁用于编码。

Transformer编码器堆叠

编码器主体由多个Transformer块组成（第39-41行）：

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)])

基础模型（mae_vit_base_patch16）配置为12层Transformer块，每层包含12个注意力头，MLP隐藏层维度为嵌入维度的4倍。

解码器详解

解码器输入构建

解码器接收编码器输出的隐藏状态，并需要恢复全部图像补丁。关键步骤是通过decoder_embed（第47行）将编码器输出映射到解码器维度，并填充掩码标记（第49行）：

self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))

在forward_decoder方法（第172行）中，掩码位置会被mask_token填充，并通过ids_restore恢复原始序列顺序。

图像重建过程

解码器输出通过decoder_pred（第58行）投影到补丁空间：

self.decoder_pred = nn.Linear(decoder_embed_dim, patch_size**2 * in_chans, bias=True)

对于16x16的补丁和RGB图像，每个补丁的输出维度为16×16×3=768。最后通过unpatchify方法（第109行）将补丁序列转换回图像格式：

def unpatchify(self, x):
    p = self.patch_embed.patch_size[0]
    h = w = int(x.shape[1]**.5)
    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

损失计算

MAE的损失函数仅计算被掩码补丁的重建误差（第198行）：

def forward_loss(self, imgs, pred, mask):
    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)  # 每个补丁的平均损失
    loss = (loss * mask).sum() / mask.sum()  # 仅计算掩码区域的损失
    return loss

norm_pix_loss参数控制是否对像素值进行归一化，实验表明这有助于稳定训练。

预定义模型配置

为方便使用，models_mae.py提供了三个预定义模型构造函数：

模型配置	函数名	补丁大小	嵌入维度	编码器深度	解码器深度
基础模型	mae_vit_base_patch16	16x16	768	12	8
大型模型	mae_vit_large_patch16	16x16	1024	24	8
巨型模型	mae_vit_huge_patch14	14x14	1280	32	8

使用示例：

model = mae_vit_base_patch16(img_size=224, in_chans=3, norm_pix_loss=True)

前向传播完整流程

forward方法（第216行）定义了完整的前向传播路径：

def forward(self, imgs, mask_ratio=0.75):
    latent, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
    pred = self.forward_decoder(latent, ids_restore)  # [N, L, p*p*3]
    loss = self.forward_loss(imgs, pred, mask)
    return loss, pred, mask

数据流向如下：