企业官网建设流程全解析

简介：一套基于CIFAR-10的MAE（Masked Autoencoder）实现，面向希望复现何凯明MAE自监督预训练效果的深度学习开发者。资源严格遵循论文细节，在CIFAR-10上完成MAE预训练，并对比了使用预训练权重微调与直接从零监督训练ViT的分类效果，用于验证自监督学习优于监督学习的数据效率结论。包内共21个文件，包含Python训练脚本（模型定义、预训练、分类器训练）、ViT-T的MAE预训练权重与监督/微调权重（.pth）、重构可视化图、README及依赖清单等，压缩包约228.74MB，目录划分清晰，便于对照实验。目前已累计1265人学习，适合已掌握基础深度学习、想通过动手实验理解MAE与自监督学习原理的研究者和工程师。从该资源中可获得完整可运行的CIFAR-10 MAE代码、预训练模型权重与实验日志，帮助快速复现论文关键结论，并在此基础上做进一步迁移或改进。

1. 在 CIFAR10 上做 MAE，先想清楚代价再动手

MAE（Masked Autoencoder）在 ImageNet 上靠 75% 掩码率逼着编码器学会全局语义，但真把它搬到 CIFAR10 上，很多人第一个版本就出现 loss 降不下去、重建图像糊成色块、TensorBoard 里曲线像心电图。原因很直接：CIFAR10 只有 32x32 分辨率，patch size 不能照搬 16，序列长度、掩码策略、学习率和权重初始化都要跟着改。这篇文章不是复述论文，而是给出一个能在单卡上跑通、能保存模型权重、能通过 TensorBoard 看到重建效果的完整实现路径。适合已经把 MAE 原理读过一遍、想在 CIFAR10 上做预训练或特征提取实验的工程师。你会得到的不是调好的黑盒，而是一套能自己改参数、排掉常见坑的工程流程。

2. MAE 的掩码重建原理与 CIFAR10 的适配参数

2.1 从图像到 token：patch embedding 的显存账

MAE 的第一步是把图像切成 patch。ImageNet 的 224x224 用 16x16 patch，得到 196 个 token；CIFAR10 的 32x32 如果也用 16x16，只剩 4 个 token，Transformer 根本学不出空间关系。常见做法是把 patch size 降到 4，即每个 patch 是 4x4x3，一张图得到 64 个 token。这个数字虽然比 ImageNet 少，但也意味着模型容量不能太大，否则单卡显存和过拟合都不好办。

在 PyTorch 里做 patch embedding，可以直接用 Conv2d 实现：stride 等于 kernel size，就是非重叠切块。以 dim=192 为例，代码如下：

import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, img_size=32, patch_size=4, in_chans=3, embed_dim=192): super().__init__() self.num_patches = (img_size // patch_size) ** 2 self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_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, device=x.device) ids_shuffle = torch.argsort(noise, dim=1) # 升序排列后，前 len_keep 个保留 ids_keep = ids_shuffle[:, :len_keep] x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).expand(-1, -1, D)) mask = torch.ones(B, L, device=x.device) mask[:, :len_keep] = 0 mask = torch.gather(mask, dim=1, index=ids_shuffle) # 恢复原始顺序，方便可视化 return x_masked, mask, ids_keep

生成 mask 时一定要把mask按照ids_shuffle重新排列，否则 TensorBoard 里画出来的掩码块和原图对不上。ids_keep是后续 decoder 恢复位置时用的索引，保存 checkpoint 时最好一并存入。

2.3 CIFAR10 版 MAE 关键超参数表

下面的参数是我在单张 24G 显存卡上跑得动的一组基线值。如果你的显存小，优先减 depth 而不是减 dim，因为 CIFAR10 的数据量小，模型太深反而学不动。

学习率换算有个简单规则：base_lr 按 batch size 256 定，你的 batch 如果是 512，lr 乘 2 而不是乘 1.5。CIFAR10 训练很快，没必要做 cosine 的重启，直接用 cosine decay 到 1e-6 就行。

3. 构建可复现的 MAE 训练代码：模型、损失函数与训练循环

3.1 编码器-解码器结构里最容易写错的维度

整个 MAE 前向过程是：patch embed 得到 token，随机掩码，编码器只处理可见 token，decoder 把可见 token 和 mask token 拼接回全量长度，最后输出每个 token 的像素值。注意 mask token 是nn.Parameter，要注册为模型参数才能被优化器更新。常见错误是把 mask token 当成固定张量，导致 decoder 永远学不到被掩码位置的信息。

下面给出一个最小的 MAE 前向逻辑，省略了 Transformer block 的具体实现，直接用nn.TransformerEncoderLayer组装：

class MAE(nn.Module): def __init__(self, img_size=32, patch_size=4, embed_dim=192, encoder_depth=8, decoder_depth=4, decoder_dim=128): 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_model=embed_dim, nhead=6, dim_feedforward=embed_dim * 4, batch_first=True, dropout=0.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_model=decoder_dim, nhead=4, dim_feedforward=decoder_dim * 4, batch_first=True, dropout=0.1) for _ in range(decoder_depth) ]) self.head = nn.Linear(decoder_dim, patch_size * patch_size * 3) def forward(self, x, mask_ratio=0.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], dim=1) 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], dim=1) x = x + self.decoder_pos # 这里需要按 ids_keep 的反操作把 token 位置还原，省略，见下节 for layer in self.decoder: x = layer(x) out = self.head(x) # (B, L+1, 48) return out, mask

编码器输出包含 cls_token 和可见 token；decoder 拼接时，也要把 cls_token 对应的位置带上。上面代码里省掉了还原顺序，实际实现要用torch.argsort(ids_keep.unsqueeze(1), dim=1)让 token 回到原始位置。

3.2 损失函数只算被掩码的部分

MAE 的 loss 只计算 mask 部分的像素重建误差，这样可以避免模型走捷径直接复制可见 patch。CIFAR10 输入已经归一化，输出层用sigmoid或者直接回归到 0~1 都行。我推荐先做归一化，再做 MSE，收敛更稳定。

def mae_loss(pred, original_patches, mask): # pred: (B, L+1, 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 的像素数会稀释 loss；CIFAR10 的 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_ratio=0.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, f'ckpt/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_step=step) writer.add_scalar('train/lr', current_lr, global_step=step) # 图片，每个 grid 放 8 张 writer.add_images('train/input', unnormalize(images[:8]), global_step=step) writer.add_images('train/masked', make_masked_image(images[:8], mask[:8]), global_step=step) writer.add_images('train/reconstruct', pred_images[:8], global_step=step)

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], dim=1) x = x + model.pos_embed for layer in model.encoder: x = layer(x) feat = x.mean(dim=1) # 全局平均池化 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))验证前向和反向都不报错，再进正式训练循环，这样最稳妥。

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