CIFAR10上MAE自监督预训练：掩码策略、PyTorch实现与调参实践

简介：基于CIFAR-10的MAE（掩码自编码器）实现包，适合深度学习研究者与有一定基础的学员复现何恺明团队的MAE自监督预训练流程。资源完整重现了论文关键设定：先用MAE对ViT进行无标签预训练，再分别微调分类器与从零监督训练，以验证自监督预训练相比纯监督学习在数据效率上的优势。包内共21个文件，以Python训练脚本、PyTorch模型权重（.pth）、README说明和TensorBoard日志为主，压缩包约228.74MB，涵盖预训练、分类训练、重构可视化等完整模块。附带模型权重、重构样例图、依赖清单及日志文件夹，便于加载后直接查看训练曲线与重构效果。目前已有1265人学习下载，适合希望快速上手MAE代码、对比自监督/监督训练结果的读者。

1. 基于 CIFAR10 的 MAE 实现：小图掩码预训练为什么要单独调

把 MAE 从 ImageNet 搬到 CIFAR10，最直接的错觉是“图变小了，训练更轻松”。实际跑一轮就会发现，32×32 输入下的 patch 数、可见 token 数、解码器负担全部变样，原版超参数直接套用大概率训不动。这个工程围绕可复现的 PyTorch 实现展开，覆盖掩码比例、编解码器维度与损失计算，再把模型权重保存和 TensorBoard 可视化同步落地。

这套实现适合想在单卡上验证自监督预训练效果的工程师，也适合准备把预训练权重迁移到分类任务的开发者。全程不需要多卡分布式，一张 12GB 显存显卡就能跑完 200 epoch。后面的代码、参数表和命令都按可直接复制执行的标准组织，并标出位置编码、掩码还原这些容易踩的点。

先明确边界：这个工程的目标是重建 32×32 图像，不是直接做分类。分类与权重复用放在最后一章的线性探针部分，会给出冻结编码器和训练分类头的完整做法。

2. CIFAR10 上 MAE 的掩码策略与编解码器架构取舍

2.1 掩码率 75% 在 32×32 输入下的可见 token 边界

MAE 的核心思路是把输入切成 patch，随机掩掉大部分，编码器只看剩余可见部分，再让解码器重建整张图。掩码本身是一种极强的数据增强——模型必须在看不到大部分像素的情况下补齐全局结构。原版在 224×224 输入、16×16 patch 下得到 196 个 token，掩码率 75% 时还有 49 个可见 token；这个数量足够编码器建立空间关系。但搬到 CIFAR10，如果继续用 patch=16，整张图只剩 4 个 token，掩掉 3 个后编码器只能看到 1 个 patch，任务直接退化成“看一个角猜整张图”。

常见做法是把 patch 尺寸缩小。CIFAR10 上两个主流选择是 patch=4 和 patch=2，token 数量和计算量差距明显：

patch=4 时 64 个 token、掩码后可见 16 个，属于“够用但紧张”的状态，训练 loss 下降速度比大图快，但也更容易提前触底。patch=2 的 token 充足，不过 256 长度的序列让自注意力计算量变成 patch=4 的 16 倍，显存和耗时都要重新评估。实际工程里一般先用 patch=4 把完整流程跑通，确认数据、损失、可视化都没问题后，再决定是否升级到 patch=2。

掩码率同样需要适配。原版 75% 在 CIFAR10 上依然可用，但如果发现重建图模糊、偏色，可以把掩码率降到 0.6 到 0.7 对比一轮。掩码率越低任务越简单，但编码器学到的高层语义会变弱，不要为了重建好看而一味降低。多组实验可以共用同一个--log-dir父目录，在 TensorBoard 里直接对比不同掩码率的 loss 曲线。

注意：mask_ratio 与 patch_size 是联动参数。patch=4 时掩码率超过 0.90，可见 token 会少于 7 个，编码器基本学不到任何空间结构。

2.2 编码器与解码器的深度、维度选型对照

MAE 的整体设计原则是“编码器重、解码器轻”。编码器只处理可见 token，实际计算量按掩码比例缩减，因此可以堆深；解码器要重建全部 token，任务简单，层数不必多，预训练结束后整个解码器直接丢弃。CIFAR10 只有 5 万张训练图，编码器堆到 ViT-Base 的 12 层很容易过拟合，常见做法是缩到 6 到 8 层。

以 patch=4、embed_dim=192 作为基线，给出三档可直接参考的配置：

解码器层数对最终重建质量的影响远小于编码器，2 到 4 层足够；编码器深度决定语义抽象能力，是迁移效果的主变量。每层 Transformer 保持标准结构：LayerNorm、多头自注意力、4 倍隐藏层 MLP、残差连接。位置编码使用学习式参数，形状固定为 (1, 64, 192)，不需要像 ViT 那样为不同分辨率做插值，因为 CIFAR10 输入尺寸恒定。

还有一个容易被忽略的点：MAE 在编码器阶段不引入任何 mask token，被掩掉的 patch 完全不参与计算；mask token 只出现在解码器输入端，用于把缺失位置补齐。这个设计让编码器的 FLOPs 和显存占用随掩码率线性下降，是 MAE 比 BEiT 类方法训练成本低的主要原因。

2.3 重建目标：逐 patch 像素归一化与 MSE 损失的配合

MAE 的监督信号是像素重建，但直接对原始像素算 MSE，模型会倾向学习全局亮度分布而忽略边缘和纹理。原版做法是对每个 patch 内的像素做独立归一化：算出该 patch 的均值和标准差，把像素标准化到零均值单位方差，解码器预测这个归一化结果；可视化时再乘回标准差、加回均值。逐 patch 归一化在小图上收益反而更大，因为 32×32 图中对象占画幅比例高，不同 patch 的亮度方差差异非常显著。对应的目标构造代码：

def patchify_normalize(x, patch_size=4): B, C, H, W = x.shape p = patch_size # 切成 (B, H/p, W/p, C, p, p) 后展平每个 patch 的像素 x = x.reshape(B, C, H // p, p, W // p, p) x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, (H // p) * (W // p), -1) mean = x.mean(dim=-1, keepdim=True) std = x.std(dim=-1, keepdim=True) + 1e-6 # 防止除零 return (x - mean) / std, mean, std

这里返回三个值：归一化目标、每个 patch 的均值、每个 patch 的标准差。后面训练循环里只把第一个值作为重建目标，后两个留给可视化反归一化用。加 1e-6 是防止纯色 patch 标准差为 0 导致除零，CIFAR10 背景区域经常出现这种情况。

损失只在被掩掉的 patch 上计算，可见 patch 不进入 loss，否则模型会集中精力“抄袭”已经看到的内容。patch=4 配置训练到 200 epoch 时，归一化 MSE 一般落到 0.03 到 0.06 区间；如果 loss 停在 0.1 以上降不动，优先检查是否漏了逐 patch 归一化，或者把可见 token 也错误地算进了 loss。

3. 用 PyTorch 搭建 CIFAR10-MAE 最小可跑通实现

3.1 数据增强与 patch embedding 的落地

MAE 预训练对数据增强的要求和其他自监督方法很不一样：不需要 MoCo 的强增强组合，也不用 SimCLR 的双分支随机裁剪。常见做法是随机裁剪加水平翻转，保留图像内容结构完整，让重建任务专注于补全掩码而不是对抗增强带来的形变。

import torch import torch.nn as nn import torchvision from torchvision import transforms def build_loader(batch_size=128, data_dir="./data"): transform = transforms.Compose([ transforms.RandomResizedCrop(32, scale=(0.7, 1.0)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2470, 0.2435, 0.2616]) ]) train_set = torchvision.datasets.CIFAR10( root=data_dir, train=True, download=True, transform=transform) return torch.utils.data.DataLoader( train_set, batch_size=batch_size, shuffle=True, num_workers=4, drop_last=True, pin_memory=True)

Normalize 的参数是 CIFAR10 预计算好的 RGB 均值与标准差，保证输入量纲稳定。需要区分的是，这里的输入标准化和 2.3 节的逐 patch 归一化是两回事：前者作用在整张图的通道维度，后者作用在损失函数内部的目标变换，两者不能混用。drop_last=True 用于保持每个 batch 形状恒定，避免最后一个不完整 batch 干扰训练。

patch embedding 用一个卷积就能实现，卷积核大小和步长都等于 patch 尺寸，输出通道等于 embed_dim，再展平成 token 序列：

class PatchEmbed(nn.Module): def __init__(self, in_channels=3, patch_size=4, embed_dim=192): super().__init__() self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): # 输入 (B, 3, 32, 32)，卷积后 (B, embed_dim, 8, 8) x = self.proj(x) B, D, H, W = x.shape return x.flatten(2).transpose(1, 2) # (B, 64, embed_dim)

flatten(2) 把 H 和 W 两个空间维度合并成 token 维度，transpose(1, 2) 把维度调整为 (B, num_tokens, D)，也就是 64 个 token、每个 192 维。token 序列之后要加上学习式位置编码，形状必须是 (1, 64, 192)，和 token 数量严格匹配，否则前向传播会直接报 shape 错误。

3.2 随机掩码采样与 ids_restore 还原逻辑

掩码采样是 MAE 里最容易写错的一步。思路是为每个样本、每个 token 生成一个独立随机数，排序后取前 len_keep 个作为可见索引；再对完整排序结果做一次逆排序，得到还原索引，供解码阶段把掩码位置补回原顺序：

def random_masking(x, mask_ratio=0.75): B, N, D = x.shape len_keep = int(N * (1 - mask_ratio)) noise = torch.rand(B, N, device=x.device) ids_shuffle = torch.argsort(noise, dim=1) # 升序排列 ids_keep = ids_shuffle[:, :len_keep] # 保留前 len_keep 个 x_masked = torch.gather( x, dim=1, index=ids_keep.unsqueeze(-1).expand(-1, -1, D)) ids_restore = torch.argsort(ids_shuffle, dim=1) # 恢复原始顺序 return x_masked, ids_keep, ids_restore

noise 是每个 token 一个 [0,1) 均匀随机数，argsort 等价于对 token 整体打乱，取前 len_keep 个就得到随机保留的子集。gather 借助 ids_keep 从完整序列里取出可见 token，作为编码器输入。ids_restore 是训练能否正确的关键：编码器输出顺序是打乱后的，解码器必须把 mask token 插回原始位置，如果顺序对不上，loss 的目标和预测会完全错位，导致 loss 曲线异常震荡但永远降不下去。

一个工程细节：掩码生成要写在模型 forward 内部，不能放进数据集。每个 epoch 都希望掩码不同，如果把掩码缓存到数据层，模型会逐渐记住固定模式，自监督预训练退化成查表。随机种子在训练脚本入口统一设置，保证单卡实验可复现；断点续训时要注意恢复后的 RNG 状态和掩码分布，必要时固定种子重新开始。

3.3 编码器-解码器前向与只算掩码位置的损失

Transformer 基础模块可以直接用 PyTorch 的nn.MultiheadAttention组合，注意输入输出都保持 (B, N, D) 的 batch_first 约定：

class TransformerBlock(nn.Module): def __init__(self, dim, num_heads): super().__init__() self.norm1 = nn.LayerNorm(dim) self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True) self.norm2 = nn.LayerNorm(dim) self.mlp = nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim)) def forward(self, x): x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] x = x + self.mlp(self.norm2(x)) return x

自注意力里 q、k、v 同源，所以传入三次相同的 norm 输出。MLP 中间层放大 4 倍是 Transformer 的惯例，激活函数用 GELU。LayerNorm 放在 attention 之前，属于 Pre-LN 结构，训练更稳定，对学习率的敏感度低。

MAE 主体把编码器、解码器串起来，前向流程包含 patch embed、位置编码、随机掩码、编码、解码、预测六个步骤：

class MaskedAutoencoder(nn.Module): def __init__(self, patch_size=4, embed_dim=192, enc_depth=6, dec_depth=4, num_heads=6, mask_ratio=0.75): super().__init__() self.mask_ratio = mask_ratio self.patch_embed = PatchEmbed(3, patch_size, embed_dim) n_tokens = (32 // patch_size) ** 2 self.pos_embed = nn.Parameter(torch.zeros(1, n_tokens, embed_dim)) self.blocks = nn.Sequential(*[ TransformerBlock(embed_dim, num_heads) for _ in range(enc_depth)]) self.norm = nn.LayerNorm(embed_dim) self.decoder_embed = nn.Linear(embed_dim, embed_dim) self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.decoder_pos_embed = nn.Parameter(torch.zeros(1, n_tokens, embed_dim)) self.decoder_blocks = nn.Sequential(*[ TransformerBlock(embed_dim, num_heads) for _ in range(dec_depth)]) self.decoder_norm = nn.LayerNorm(embed_dim) patch_area = patch_size * patch_size self.decoder_pred = nn.Linear(embed_dim, patch_area * 3) def forward(self, x): x = self.patch_embed(x) + self.pos_embed x_masked, ids_keep, ids_restore = random_masking(x, self.mask_ratio) x = self.blocks(x_masked) x = self.norm(x) # 编码器输出 (B, len_keep, D) x = self.decoder_embed(x) n_mask = ids_restore.shape[1] - x.shape[1] mask_tokens = self.mask_token.expand(x.shape[0], n_mask, -1) x = torch.cat([x, mask_tokens], dim=1) x = torch.gather(x, dim=1, index=ids_restore.unsqueeze(-1).expand(-1, -1, x.shape[2])) x = x + self.decoder_pos_embed x = self.decoder_blocks(x) x = self.decoder_norm(x) pred = self.decoder_pred(x) # (B, N, patch_area*3) return pred, ids_keep

decoder_pred 的输出维度是 patch_area×3，即每个 token 重建一个 patch 的 RGB 像素。位置编码在掩码前就已经加入，保证可见 token 携带位置信息；解码器的位置编码单独一套，加在补齐后的完整序列上。

损失函数按掩码位置求平均，避免可见 token 稀释监督信号：

def mae_loss(pred, target, ids_keep): B, N, _ = pred.shape mask = torch.ones(B, N, device=pred.device) for b in range(B): mask[b, ids_keep[b]] = 0 # 可见 token 不计入 loss loss = ((pred - target) ** 2).mean(dim=-1) return (loss * mask).sum() / mask.sum()

训练循环把三块拼起来：前向拿到预测和 ids_keep，用 2.3 节的 patchify_normalize 构造目标，最后套 mae_loss：

optimizer = torch.optim.AdamW(model.parameters(), lr=1.5e-4, weight_decay=0.05) for epoch in range(200): for img, _ in loader: img = img.cuda() pred, ids_keep = model(img) target, _, _ = patchify_normalize(img, patch_size=4) loss = mae_loss(pred, target, ids_keep) optimizer.zero_grad() loss.backward() optimizer.step()

4. 训练配置、模型权重保存与 TensorBoard 可视化

4.1 训练超参数表与 warmup+cosine 调度

MAE 训练最忌讳直接照搬 ImageNet 的 1600 epoch 配置。CIFAR10 只有 5 万张训练图，200 epoch 已经足够，再多容易过拟合到训练集的 RGB 分布上。优化器选 AdamW，权重衰减 0.05；基础学习率按 batch size 线性缩放，公式是 lr = base_lr × batch_size / 256。以下是 patch=4、embed_dim=192、单卡 12GB 显存下的基线配置：

warmup 在自监督训练里不是可选项。训练初期梯度方向噪声大，位置编码和 mask token 都是新初始化的参数，直接上大学习率会让它们剧烈震荡，后期再难收敛。cosine 调度让后半程逐步降低学习率，配合逐 patch 归一化的 MSE 目标，可以找到更平滑的极小值。

一个常见坑是 weight decay 误伤位置编码和 mask token。这两类参数不存在“权重越大越复杂”的问题，应该从 weight decay 里排除。实现上把参数分成 no_decay 和 decay 两组，前者包含 pos_embed、mask_token 以及所有 bias 和 LayerNorm 参数。否则训练后期会出现重建图像整体变灰、高频细节丢失的现象。

4.2 checkpoint 的保存策略与严格/非严格加载

模型权重是工程交付的一部分，保存策略直接决定后续微调是否顺利。常见做法是每个 epoch 结束后在验证集上算一次重建 loss，只保留最优和最近两个 checkpoint，避免 200 个 epoch 把磁盘占满。保存内容要覆盖模型、优化器、调度器和 epoch 计数，缺了任何一项都无法干净地断点续训：

def save_checkpoint(model, optimizer, scheduler, epoch, path): torch.save({ "epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), "mask_ratio": model.mask_ratio, }, path)

加载分两种场景。继续训练要求 key 完全匹配，直接load_state_dict(torch.load(path)["model"])；微调分类任务时只需要编码器，要用 strict=False 加载，并主动过滤掉 decoder 相关权重：

ckpt = torch.load("mae_cifar10_best.pth")["model"] new_state = {k: v for k, v in ckpt.items() if k.startswith("blocks") or k.startswith("pos_embed")} model.load_state_dict(new_state, strict=False)

过滤条件保留 blocks 和输入侧位置编码 pos_embed，丢弃所有 decoder_ 开头的权重向量。之所以“宁多勿少”，是因为 strict=False 会静默跳过所有缺失的 key，漏掉 pos_embed 不会报错，但线性探针精度会明显下降。建议保存 checkpoint 时顺带存一份参数配置字典，加载时先对比 patch_size、embed_dim 是否一致，避免不同架构的权重硬载导致静默错误。

4.3 TensorBoard 标量、权重直方图与重建图像

TensorBoard 可视化分三层：标量曲线、权重直方图、图像。标量至少记录训练 loss、验证 loss、学习率三项；直方图记录编码器第一个 block 的 attention 权重和 mask token 数值分布，用于观察训练是否饱和；图像部分把原始图、掩码图、重建图拼成一个 grid，是最直观的调试入口：

from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter(log_dir="runs/mae_cifar10") writer.add_scalar("train/loss", loss.item(), global_step=epoch) writer.add_scalar("train/lr", optimizer.param_groups[0]["lr"], global_step=epoch) writer.add_histogram("encoder/attn_weight", attn_weight.detach(), global_step=epoch) writer.add_images("vis/reconstruction", grid, global_step=epoch)

grid 的构造方式是把掩码图和重建图按 batch 维度拼成 (2B, 3, 32, 32)，add_images 会自动排成网格展示。重建图反归一化是这里最容易翻车的地方：模型预测的是逐 patch 归一化后的值，必须先乘以 2.3 节保存的 std、加上 mean，再做数据集 Normalize 的逆变换，才能得到真实 RGB。漏掉任何一步，TensorBoard 里看到的就是一片灰噪。

提示：add_images 显示全灰时，优先检查反归一化链路，而不是怀疑模型没收敛。用 20 epoch 左右的中间 checkpoint 先验证可视化管线，能节省大量排错时间。

4.4 一条命令启动训练并打开 tensorboard

训练入口用 argparse 把 patch_size、embed_dim、mask_ratio、epochs、log_dir 全部参数化，方便掩码率、深度等消融实验的对比。启动训练和 TensorBoard 服务分别开两个终端：

python train_mae.py \ --patch-size 4 --embed-dim 192 --enc-depth 6 --dec-depth 4 \ --mask-ratio 0.75 --epochs 200 --batch-size 128 \ --data-dir ./data --log-dir runs/mae_cifar10 tensorboard --logdir runs/mae_cifar10 --port 6006

tensorboard 命令默认递归扫描 logdir 下的所有子目录，--port 指定端口避免和已有进程冲突。浏览器打开 localhost:6006 后，SCALARS 页看 loss 和学习率曲线，IMAGES 页看重建效果。跑了多组 mask_ratio 对比实验时，把 log_dir 统一成 runs 下的不同子目录，再用tensorboard --logdir runs加载父目录，TensorBoard 会自动按子目录分组展示对比曲线。

5. 重建质量验证与权重复用技巧

5.1 PSNR 与 SSIM 评估重建结果

loss 数值只能反映训练进度，不能说明重建质量的观感。补两个图像级指标：PSNR 衡量像素误差，SSIM 衡量结构相似性。PSNR 可以手写，SSIM 直接用 skimage 或 torcheval 提供：

import torch.nn.functional as F def psnr(pred, target, max_val=1.0): mse = F.mse_loss(pred, target) return 10 * torch.log10(max_val * max_val / mse)

评估时选取验证集上随机 32 张图，逐张计算后取平均。CIFAR10 基线配置的 PSNR 一般在 26 到 29dB 之间；低于 24dB 说明重建基本不可辨，优先检查掩码还原顺序和损失是否正确。SSIM 建议按 patch 计算而不是整图，因为 32×32 整图 SSIM 对边缘位移过于敏感。

5.2 冻结编码器做线性探针微调

权重复用的起点是线性探针：冻结完整编码器，接全局平均池化和一个 10 类分类头，只训练分类头。这个实验能快速判断预训练质量，超参数也少：

model.requires_grad_(False) head = nn.Linear(192, 10) optimizer = torch.optim.AdamW(head.parameters(), lr=3e-3) # 训练 30 epoch，其余数据增强与 3.1 节一致

线性探针只更新分类头，梯度不会回流到编码器，所以可以用较大的 3e-3 学习率。CIFAR10 上线性探针精度达到 60% 说明 MAE 预训练有效，超过 75% 说明掩码比例和训练时长已经调得比较到位。如果精度远低于 60%，优先怀疑训练不充分或掩码率太高，而不是架构问题。

5.3 TensorBoard 上判断训练正常的三个信号

第一个信号是训练 loss 与验证 loss 同步下降，中间没有突然尖峰；尖峰通常对应学习率设置不当或某个 batch 出现 NaN 梯度。第二个信号是 encoder 输出特征的标准差稳定在合理区间，不持续增大也不归零；持续增大说明位置编码和权重 decay 参数组没有分开。第三个信号是重建图从 20 epoch 左右开始能看出物体轮廓，60 epoch 后边缘变清晰；如果 60 epoch 还在输出平均色块，回头看 4.1 节的超参数表逐项核对。三者同时满足，checkpoint 就可以放心用于后续任务。

本文还有配套的精品资源，点击获取