简介一套基于CIFAR-10的MAEMasked Autoencoder实现面向希望复现何凯明MAE自监督预训练效果的深度学习开发者。资源严格遵循论文细节在CIFAR-10上完成MAE预训练并对比了使用预训练权重微调与直接从零监督训练ViT的分类效果用于验证自监督学习优于监督学习的数据效率结论。包内共21个文件包含Python训练脚本模型定义、预训练、分类器训练、ViT-T的MAE预训练权重与监督/微调权重.pth、重构可视化图、README及依赖清单等压缩包约228.74MB目录划分清晰便于对照实验。目前已累计1265人学习适合已掌握基础深度学习、想通过动手实验理解MAE与自监督学习原理的研究者和工程师。从该资源中可获得完整可运行的CIFAR-10 MAE代码、预训练模型权重与实验日志帮助快速复现论文关键结论并在此基础上做进一步迁移或改进。1. 在 CIFAR10 上做 MAE先想清楚代价再动手MAEMasked Autoencoder在 ImageNet 上靠 75% 掩码率逼着编码器学会全局语义但真把它搬到 CIFAR10 上很多人第一个版本就出现 loss 降不下去、重建图像糊成色块、TensorBoard 里曲线像心电图。原因很直接CIFAR10 只有 32x32 分辨率patch size 不能照搬 16序列长度、掩码策略、学习率和权重初始化都要跟着改。这篇文章不是复述论文而是给出一个能在单卡上跑通、能保存模型权重、能通过 TensorBoard 看到重建效果的完整实现路径。适合已经把 MAE 原理读过一遍、想在 CIFAR10 上做预训练或特征提取实验的工程师。你会得到的不是调好的黑盒而是一套能自己改参数、排掉常见坑的工程流程。2. MAE 的掩码重建原理与 CIFAR10 的适配参数2.1 从图像到 tokenpatch embedding 的显存账MAE 的第一步是把图像切成 patch。ImageNet 的 224x224 用 16x16 patch得到 196 个 tokenCIFAR10 的 32x32 如果也用 16x16只剩 4 个 tokenTransformer 根本学不出空间关系。常见做法是把 patch size 降到 4即每个 patch 是 4x4x3一张图得到 64 个 token。这个数字虽然比 ImageNet 少但也意味着模型容量不能太大否则单卡显存和过拟合都不好办。在 PyTorch 里做 patch embedding可以直接用 Conv2d 实现stride 等于 kernel size就是非重叠切块。以 dim192 为例代码如下import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, img_size32, patch_size4, in_chans3, embed_dim192): super().__init__() self.num_patches (img_size // patch_size) ** 2 self.proj nn.Conv2d(in_chans, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): # x: (B, 3, 32, 32) x self.proj(x) # (B, embed_dim, 8, 8) x x.flatten(2).transpose(1, 2) # (B, 64, embed_dim) return x这里flatten(2)把 8x8 的空间维度摊平成 64 个 token再transpose得到(B, 64, C)。注意不要再在 p 后面拼一个Linear因为 Conv2d 已经做了线性映射叠加反而增加无谓参数。2.2 掩码生成为什么 75% 掩码率在 CIFAR10 上需要回调MAE 论文默认 75% 掩码对于 196 个 token 的 ImageNet 足够困难但 CIFAR10 只有 64 个 token如果也砍掉 75%编码器只看 16 个 token对一张 32x32 的小图来说信息残留太少预训练早期容易崩塌。我习惯在 CIFAR10 上先跑 60% 到 65% 掩码率等重建 loss 稳定后再调高到 75% 做对比实验。掩码生成有两种方式一种是随机采样 int 索引一种是用 Bernoulli 分布生成 mask。后者代码更短且容易控制掩码比例def random_masking(x, mask_ratio): B, L, D x.shape # x: (B, seq_len, dim) len_keep int(L * (1 - mask_ratio)) noise torch.rand(B, L, devicex.device) ids_shuffle torch.argsort(noise, dim1) # 升序排列后前 len_keep 个保留 ids_keep ids_shuffle[:, :len_keep] x_masked torch.gather(x, dim1, indexids_keep.unsqueeze(-1).expand(-1, -1, D)) mask torch.ones(B, L, devicex.device) mask[:, :len_keep] 0 mask torch.gather(mask, dim1, indexids_shuffle) # 恢复原始顺序方便可视化 return x_masked, mask, ids_keep生成 mask 时一定要把mask按照ids_shuffle重新排列否则 TensorBoard 里画出来的掩码块和原图对不上。ids_keep是后续 decoder 恢复位置时用的索引保存 checkpoint 时最好一并存入。2.3 CIFAR10 版 MAE 关键超参数表下面的参数是我在单张 24G 显存卡上跑得动的一组基线值。如果你的显存小优先减 depth 而不是减 dim因为 CIFAR10 的数据量小模型太深反而学不动。参数取值说明patch_size432x32 切成 8x8 网格共 64 tokenmask_ratio0.6~0.75推荐 0.65 起步encoder_depth8每层是标准 Transformer blockencoder_dim192比 ImageNet 版小一半decoder_depth4decoder 不用太深decoder_dim128与 decoder 输出通道匹配lr1.5e-4batch 256 时的基准值按线性缩放weight_decay0.05只对非 bias 和 norm 参数生效warmup_epochs5前 5 个 epoch 线性升 lr学习率换算有个简单规则base_lr 按 batch size 256 定你的 batch 如果是 512lr 乘 2 而不是乘 1.5。CIFAR10 训练很快没必要做 cosine 的重启直接用 cosine decay 到 1e-6 就行。3. 构建可复现的 MAE 训练代码模型、损失函数与训练循环3.1 编码器-解码器结构里最容易写错的维度整个 MAE 前向过程是patch embed 得到 token随机掩码编码器只处理可见 tokendecoder 把可见 token 和 mask token 拼接回全量长度最后输出每个 token 的像素值。注意 mask token 是nn.Parameter要注册为模型参数才能被优化器更新。常见错误是把 mask token 当成固定张量导致 decoder 永远学不到被掩码位置的信息。下面给出一个最小的 MAE 前向逻辑省略了 Transformer block 的具体实现直接用nn.TransformerEncoderLayer组装class MAE(nn.Module): def __init__(self, img_size32, patch_size4, embed_dim192, encoder_depth8, decoder_depth4, decoder_dim128): super().__init__() self.patch_embed PatchEmbed(img_size, patch_size, 3, embed_dim) self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed nn.Parameter( torch.zeros(1, self.patch_embed.num_patches 1, embed_dim)) self.encoder nn.ModuleList([ nn.TransformerEncoderLayer( d_modelembed_dim, nhead6, dim_feedforwardembed_dim * 4, batch_firstTrue, dropout0.1) for _ in range(encoder_depth) ]) self.decoder_embed nn.Linear(embed_dim, decoder_dim) self.mask_token nn.Parameter(torch.zeros(1, 1, decoder_dim)) self.decoder_pos nn.Parameter( torch.zeros(1, self.patch_embed.num_patches 1, decoder_dim)) self.decoder nn.ModuleList([ nn.TransformerEncoderLayer( d_modeldecoder_dim, nhead4, dim_feedforwarddecoder_dim * 4, batch_firstTrue, dropout0.1) for _ in range(decoder_depth) ]) self.head nn.Linear(decoder_dim, patch_size * patch_size * 3) def forward(self, x, mask_ratio0.65): x self.patch_embed(x) # (B, 64, C) B, L, _ x.shape # 拼接 cls_token这个 token 永远不被掩码 x torch.cat([self.cls_token.expand(B, -1, -1), x], dim1) x x self.pos_embed x_visible, mask, ids_keep random_masking(x, mask_ratio) for layer in self.encoder: x_visible layer(x_visible) # decoder 侧要还原顺序 x_full self.decoder_embed(x_visible) mask_tokens self.mask_token.repeat(B, L 1 - x_visible.shape[1], 1) x torch.cat([x_full, mask_tokens], dim1) x x self.decoder_pos # 这里需要按 ids_keep 的反操作把 token 位置还原省略见下节 for layer in self.decoder: x layer(x) out self.head(x) # (B, L1, 48) return out, mask编码器输出包含 cls_token 和可见 tokendecoder 拼接时也要把 cls_token 对应的位置带上。上面代码里省掉了还原顺序实际实现要用torch.argsort(ids_keep.unsqueeze(1), dim1)让 token 回到原始位置。3.2 损失函数只算被掩码的部分MAE 的 loss 只计算 mask 部分的像素重建误差这样可以避免模型走捷径直接复制可见 patch。CIFAR10 输入已经归一化输出层用sigmoid或者直接回归到 0~1 都行。我推荐先做归一化再做 MSE收敛更稳定。def mae_loss(pred, original_patches, mask): # pred: (B, L1, p*p*3), mask: (B, L) 注意 pred 包含 cls 位 pred pred[:, 1:, :] # 去掉 cls token target original_patches[:, :, :] # (B, L, p*p*3) loss (pred - target) ** 2 loss loss.mean(dim-1) # 每个 token 内取平均 mask mask.bool() # True 表示被掩码 loss (loss * mask).sum() / mask.sum() return loss为什么要在 token 内先取平均再乘 mask因为如果你先算整张图 MSE 再乘 mask不同 patch 的像素数会稀释 lossCIFAR10 的 patch 是 4x4总共 48 像素数量不大但也要保持一致性。target 要从原始图像切出来不能复用random_masking的输出因为那已经是可见 token 了。3.3 训练循环里的三个隐藏细节训练循环最简单的版本谁都写得出来但直接跑你会遇到三个问题一是不给 warmup早期 loss 直接 inf二是没有按 step 更新 TensorBoard曲线间隔不均匀三是不做梯度裁剪偶发 NaN 后模型权重就废了。我一般在train_one_epoch里用tqdm显示当前 loss每 50 步写一次 TensorBoard并顺手做梯度裁剪optimizer.zero_grad() output, mask model(images, mask_ratio0.65) loss mae_loss(output, patch_target(images), mask) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() # 如果按 step 调整参数说明clip_grad_norm_的 max_norm 设为 1.0 是针对 CIFAR10 小模型的经验值大模型可以放到 3.0。scheduler.step()如果你的 scheduler 构建时用的是optimizer那么要在optimizer.step()之后调用如果你用 epoch 级的 cosine 退火就放到每个 epoch 末尾。4. 模型权重的保存/加载与 TensorBoard 可视化落地4.1 checkpoint 里不应该只存 state_dict训练 MAE 这种自监督模型最忌讳只存model.state_dict()。因为后续做线性评估、微调或者改掩码率继续训练都需要知道之前的训练配置。我至少会存五样东西模型权重、优化器状态、当前 epoch、mask_ratio、best loss。checkpoint { epoch: epoch, model: model.state_dict(), optimizer: optimizer.state_dict(), mask_ratio: args.mask_ratio, best_loss: best_loss, } torch.save(checkpoint, fckpt/mae_cifar10_{epoch:02d}.pth)加载的时候注意如果你只加载模型权重去测试不要直接load_state_dict就完事。要先构建相同结构的 MAE 实例再load_state_dict并且如果模型里用了nn.TransformerEncoderLayer不同 PyTorch 版本之间可能有layer_norm_eps的差异最好从 checkpoint 里读一下args确保结构一致。4.2 TensorBoard 记录重建图像与原图对比TensorBoard 的可视化分两类标量曲线和图片。标量部分记录 total loss、重建 loss 和 lr 就足够。图片部分要记录三个图原图、带掩码的图、重建图。掩码图可以用 mask 直接乘到 patch 上令被掩码区域变成灰色或者直接 0。from torch.utils.tensorboard import SummaryWriter writer SummaryWriter(runs/mae_cifar10) # 标量 writer.add_scalar(train/loss, loss.item(), global_stepstep) writer.add_scalar(train/lr, current_lr, global_stepstep) # 图片每个 grid 放 8 张 writer.add_images(train/input, unnormalize(images[:8]), global_stepstep) writer.add_images(train/masked, make_masked_image(images[:8], mask[:8]), global_stepstep) writer.add_images(train/reconstruct, pred_images[:8], global_stepstep)make_masked_image的实现逻辑是把每个 patch 对应的 mask 值乘到该 patch 的所有像素上mask 为 1 的位置置 0。如果只想看一块区域还可以叠加一个torch.where把 mask 区域替换成灰色0.5。4.3 从曲线判断训练是否健康TensorBoard 里最容易误判的是 loss 曲线。MAE 的 loss 一开始会快速下降随后进入平台期这是正常的。你要关注的不是曲线本身而是两个异常模式第一个是 loss 在某个 step 突然跳高后恢复这通常是学习率步长没对齐 scheduler第二个是 loss 持续不降维持在一个固定值附近这时 90% 的情况是 mask 生成逻辑里ids_shuffle没有做反序还原导致 decoder 永远在乱序 token 上做回归。判断方法很简单在验证集上固定一个 batch可视化重建图和原图的 SSE 热力图如果错误集中在图像边缘说明位置编码没学好。4.4 用模型权重可视化注意力分布保存好 checkpoint 之后不着急立刻做下游任务可以先看一下编码器的注意力模式。对 CIFAR10 这种小图头几个 Transformer block 的 attention 通常集中在 patch 自身越到深层越会关注远距离 patch。写一个 hook 把所有 attention 矩阵抓出来attention_weights {} def hook_fn(name): def fn(module, inp, out): # 需要从 nn.MultiheadAttention 的输出去拿 attn可注册 forward hook attention_weights[name] out[1] if isinstance(out, tuple) else inp return fn然后对第 4 个 encoder block 的 attention 做平均投影到 8x8 的 patch 网格上用writer.add_image画出来。这个图和重建图对照看你能直观判断模型是学到了纹理还是学到了轮廓。如果注意力全部集中在局部 3x3 邻域说明 encoder 深度设太多了CIFAR10 并不需要 8 层。5. 用 TensorBoard 验证表征质量重建之外要看线性探针重建 loss 降到 0.1 不代表特征有用因为 CIFAR10 的像素空间很小模型可以直接记住像素分布。我建议在最后一章落地一个快速验证技巧冻结编码器只训练一个全局池化加分类头在 CIFAR10 测试集上看 top-1。具体做法是把 MAE 拆成encoder_without_head对每张图输出 64 个 patch token再做一个可学习的平均池化然后过一个nn.Linear(192, 10)。线性探针代码很短model.eval() feats [] with torch.no_grad(): for images in tqdm(test_loader): images images.cuda() x model.patch_embed(images) x torch.cat([model.cls_token.expand(x.shape[0], -1, -1), x], dim1) x x model.pos_embed for layer in model.encoder: x layer(x) feat x.mean(dim1) # 全局平均池化 feats.append(feat) feats torch.cat(feats) # 用 LogisticRegression 训练分类头比如 sklearn 的这个验证的意义在于TensorBoard 里看到的重建图只是自监督任务的代理指标线性准确率才是表征质量的近似度量。CIFAR10 上随机初始化特征做线性探针大约 10% 到 20%训练 100 个 epoch 的 MAE 特征一般能到 60% 以上。如果你发现重建图很清晰但线性准确率不足 50%优先检查是否用了数据增强。MAE 预训练阶段只做随机翻转和 cutout 就够不要加 mixup因为混合图像的 patch token 会让掩码重建任务变得不一致。最后的调试技巧是在 TensorBoard 的graph标签页里看一眼模型结构确认 MaskToken 和 PositionEmbedding 都挂在了正确模块下。很多隐性 bug 是 checkpoint 里缺了mask_token参数导致加载后重建图出现周期性条纹。用torch.save前执行一次model(torch.randn(4,3,32,32))验证前向和反向都不报错再进正式训练循环这样最稳妥。本文还有配套的精品资源点击获取