自监督学习新范式:MoBY在Transformer上超越MoCo v3的实践解析
2026/9/16 2:36:42 网站建设 项目流程

简介:这一深度学习源码包聚焦自监督学习方法MoBY,面向具备一定深度学习基础的研究者与算法工程师,用于理解并复现将MoCo v2与BYOL优势结合的新型训练范式。方法以视觉变换器为主干网络,经过三百轮训练即可在图像分类标准数据集上取得较高线性评估准确率:使用较小规模的DeiT模型达到七十二点八,使用Swin模型达到七十五点零,且相较于同类方法MoCo v3与DINO,技巧更轻、复现成本更低。包内共三十三个文件,以十八个Python源码为主,包含训练主程序、模型搭建、数据加载与学习率调度等模块,另附九个配置文件、两个说明文档和两张效果示意图,整体压缩包大小仅一点一一兆字节,目录划分清晰,便于按需查阅。目前已有五百一十六人学习下载,对希望对比主流自监督学习算法或迁移MoBY到自身任务的中高级学习者,具有直接参考价值。代码还附带了线性评估脚本与历史配置存档,可辅助读者快速验证模型性能、复现不同实验设置,从而更深入地掌握这一兼顾精度与简洁性的自监督学习方案。

1. 自监督学习新范式:MoBY 如何在Transformer上超越MoCo v3

自监督学习过去两年在视觉领域的进展,基本绕不开两个名字:MoCo 和 BYOL。前者的对比学习思路依赖负样本构造,后者则用不对称的两个网络分支避免了负样本依赖。MoBY 这个名字在社区里不算高调,但它把这两条路线压进了同一套框架,直接用 Vision Transformer(ViT)家族做主干,在 ImageNet-1K 上只训 300 epoch,线性评估就在 DeiT-S 和 Swin-T 上分别拿到 72.8% 和 75.0% 的 top-1 准确率。对比同期的 MoCo v3 和 DINO,MoBY 的 trick 更轻,训练也更稳定,是理解自监督学习从 ResNet 迁移到 Transformer 的一个干净样本。这份源码包完整覆盖了数据加载、模型构建、训练器和线性评估脚本,适合正在做自监督实验、需要快速跑通一套对比基线的从业者。

2. 机制与目录结构:MoCo v2 和 BYOL 是怎么在 MoBY 里共存的

2.1 从两个流派的核心矛盾说起

MoCo v2 的核心是动态字典和动量更新。它维护一个队列作为负样本池,key 编码器的参数通过指数移动平均(EMA)从 query 编码器复制,这样在训练过程中负样本既能覆盖足够多的类别,又不会因为参数更新太快而让对比任务失去难度。

BYOL 走的是另一条路。它没有负样本,靠的是两个结构相同但参数不同步的编码器,加上一个额外的 predictor 结构。目标网络不直接吃梯度,而是用动量方式缓慢逼近在线网络,所以模型学到的表征不会崩坏,也就是社区常说的"免负样本对比"。

MoBY 的贡献在于,它没有把这两个框架简单地拼在一起,而是在损失函数、梯度更新策略和主干结构上做了统一的再设计。从源码包里的moby_main.py能清楚看到,train 状态下的 forward 流程包含两条数据增强分支,一条走 student(online)路径,一条走 teacher(target)路径,同时在内存里维护一个尺寸可配置的队列。这实际上是"MoCo 的队列 + BYOL 的跨视图预测"的组合,而非简单的孪生网络。

2.2 源码包目录的功能划分

拿到 zip 解压后,第一阶段要先把目录结构和职责摸清楚,否则后面改参数会到处找文件。源码包根目录是这样的:

Transformer-SSL-main/ ├── moby_main.py # 训练入口,支持多卡分布式 ├── moby_linear.py # 线性评估脚本 ├── get_started.md # 快速上手说明 ├── utils.py # 分布式、日志、统计工具 ├── optimizer.py # AdamW + Layer-wise lr decay ├── config.py # 全局配置解析 ├── lr_scheduler.py # 余弦退火调度 ├── logger.py # 日志记录 ├── zipreader.py # zip 格式数据集读取 ├── samplers.py # 分布式采样器 ├── build.py # 数据加载器工厂 ├── custom_image_folder.py # 自定义 ImageFolder 实现 ├── cached_image_folder.py # 带缓存的 ImageFolder ├── models/ │ ├── moby.py # MoBY 核心类 │ ├── swin_transformer.py # Swin-T 实现 │ ├── __init__.py # 模型注册表 ├── configs/ │ ├── moby_swin_tiny_bn.yaml │ ├── moby_swin_tiny.yaml │ └── moby_deit_small.yaml └── figures/

这个布局对训练类项目来说已经算克制,没有把数据增强、调度器、分布式逻辑全部堆进 main。关键文件是models/moby.pymoby_main.py,前者定义模型的 forward 和 loss 计算,后者负责控制训练循环和 checkpoint 保存。configs/下三个 yaml 对应三种主干:Swin-T 带 BatchNorm、Swin-T 不带 BN、DeiT-S,三者性能差异在最后做实验对比时非常明显。

2.3 为什么不直接用现成的 MoCo v3

做 Transformer 自监督时,MoCo v3 在某些 batch size 下会出现训练崩溃,本质原因是 BatchNorm 在 ViT 上会导致梯度传播不稳定。MoBY 的策略是让 momentum encoder(teacher 分支)完全不做梯度回传,只在每次迭代后用 student 的参数做动量更新,并在 Swin-T 上默认启用 BatchNorm 来稳定训练。

用一句话概括:MoCo v3 的 trick 大多花在"如何在 ViT 上让对比学习稳定",MoBY 则是把 BYOL 的稳定性吸收进来,让 MoCo 的队列不必承担稳定性压力。这也是为什么在同样 300 epoch 下,MoBY 的收敛曲线比 MoCo v3 要平缓得多。

提示:源码包里models/moby.pymomentum参数默认设的是 0.99,如果数据量小或者 batch size 只有 256,建议先降到 0.98 观察 loss 曲线,防止 momentum 过大导致 teacher 更新过慢。

2.4 损失函数的代码对照

models/moby.py里最关键的一段,是 loss 计算的部分。核心思路是,把 student 特征和被 queue 裁剪过的 teacher 特征做 InfoNCE 损失,同时保留 queue 的先进先出更新逻辑。简化后的伪代码如下:

# 关键实现来自 models/moby.py,已删除 BN 和 predictor 细节 def forward(self, img_q, img_k, queue, queue_ptr): # student 分支:吃第一个视图 q = self.student_encoder(img_q) q = nn.functional.normalize(q, dim=-1) # teacher 分支:吃第二个视图,不计算梯度 with torch.no_grad(): k = self.teacher_encoder(img_k) k = nn.functional.normalize(k, dim=-1) # 从 queue 中取出一批旧特征作为负样本 l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1) l_neg = torch.einsum('nc,ck->nk', [q, queue.clone().detach()]) # InfoNCE 损失 logits = torch.cat([l_pos, l_neg], dim=1) labels = torch.zeros(logits.shape[0], dtype=torch.long) loss = nn.CrossEntropyLoss()(logits / self.temperature, labels) # 更新队列:当前 batch 的 teacher 特征入队 ptr = int(queue_ptr) queue[:, ptr:ptr + batch_size] = k.T ptr = (ptr + batch_size) % queue_size return loss, q, k

这段代码看着简单,但有三个细节值得注意。第一,l_neg的计算用了torch.einsum('nc,ck->nk'),这里的queue是转置过的,形状为[dim, queue_size],所以在做矩阵乘法之前不需要调换维度。第二,queue[:, ptr:ptr + batch_size] = k.T用的是队列指针而非直接拼接,这样整个队列的内存是固定的,不会随训练长度增长。第三,损失函数用的是普通交叉熵,因为l_pos永远是排序第 0 的样本,所以labels全部置 0 即可。

注意:如果你的实验里 batch size 特别大,队列指针的更新要放在 loss 计算之后,否则当前 batch 的 k 会把自己当作负样本参与计算。

3. 数据链路实战:从 zip 数据集到分布式采样

3.1 ImageFolder 的扩展:custom_image_folder.py 在解决什么问题

自监督训练的数据读取和普通监督分类有个关键区别:每个样本要读两次,而且两次读取要用不同的随机增强。custom_image_folder.py在 PyTorch 原生ImageFolder的基础上做了两个调整,一个是允许同一个样本在__getitem__里返回两条增强分支,另一个是支持 zip 包内直接读取,避免解压上万张小图时 inode 耗尽。

class CustomImageFolder(datasets.ImageFolder): def __init__(self, root, transform=None, two_crop=False): super().__init__(root, transform=transform) self.two_crop = two_crop def __getitem__(self, index): path, target = self.samples[index] img = self.loader(path) if self.two_crop is False: return self.transform(img), target # 这里是 MoBY 自定义的核心: # 同一张图走两次 transform,生成两个视角 view1 = self.transform(img) view2 = self.transform(img) return view1, view2

two_crop=True时,每次取样本都会在同一张原图上应用两次随机增强。这个设计直接对应 MoCo 系方法里的 query 和 key 两个视角。实际使用中,第二次增强的强度不能和第一次完全一样,否则模型会通过"记住像素噪声"来走捷径,常见的做法是把 color jitter 的强度在第二个分支上再乘一个系数。

3.2 zipreader.py:省掉解压环节的 IO 优化

zipreader.py的存在让数据可以从 zip 包内直接读取,不必把所有图片解压到磁盘。这一点在 ImageNet-1K 这种百万级图片数据集上非常有用,因为解压后的文件数量会严重影响文件系统的 inode 使用。

class ZipReader: def __init__(self, zip_path): self.zip_path = zip_path self.zip_file = zipfile.ZipFile(zip_path) # 建立文件名到 zip 内部位置的索引 self.name_to_info = { info.filename: info for info in self.zip_file.infolist() } def read(self, name): info = self.name_to_info[name] return self.zip_file.read(info)

用 zip 格式做训练数据容器,最大的收益是随机读取的局部性更好,SSD 在顺序读 zip 内的连续文件段时吞吐量远高于随机读百万个小文件。但代价是 zip 的压缩率如果设置过高,CPU 解压会成为瓶颈,建议打包时用ZIP_STORED模式,只做容器不打压缩。

3.3 samplers.py 与分布式采样的坑

多卡训练时,每张卡负责的数据子集不能重叠,同时为了保证每个 epoch 的 shuffle 质量,samplers.py里实现对分布式环境做了兼容。用DistributedSampler的常见问题有两个,一是没有在 epoch 开始时调用set_epoch,导致每个 epoch 的 shuffle 顺序完全一样;二是 drop_last 参数没设置,导致最后不完整的 batch 在 BN 统计上产生偏差。

moby_main.py的训练循环可以看出,作者在train_one_epoch里显式调用了sampler.set_epoch(epoch),这一步不能漏,否则模型的泛化会下降。对于 batch size 小于 512 的小实验,建议 drop_last 开成 True,因为 MoCo 系的损失对负样本数量比较敏感,最后几个不完整 batch 会让损失计算出现波动。

4. 训练配置与优化细节:三个 yaml 的横向对比

4.1 配置文件的模块化设计

configs/下三个 yaml 文件分别对应 DeiT-S、Swin-T 带 BN、Swin-T 不带 BN。这种按主干和归一化方式拆分的做法,让每个实验的复现成本降到最低。先看moby_swin_tiny.yaml的核心字段:

# configs/moby_swin_tiny.yaml model: arch: 'moby_swin_tiny' pretrained: false use_bn: false # Swin-T 不启用 BatchNorm momentum: 0.99 # teacher 分支的移动平均系数 train: batch_size: 256 # 单卡 batch size,8 卡就是 2048 base_lr: 2.5e-4 warmup_epochs: 20 epochs: 300 weight_decay: 0.05 layer_decay: 0.65 # Transformer 层间学习率衰减系数 data: data_path: '/path/to/imagenet' zip_mode: true num_workers: 12

batch_sizebase_lr是联动的关系。MoBY 默认的base_lr是按总 batch size 2048 计算的,如果你只有单卡且 batch size 只有 128,学习率要相应缩到2.5e-4 * 128 / 2048 ≈ 1.56e-5才能保证收敛行为一致。另外layer_decay这个参数是从 BEiT 那边借来的,核心思想是浅层学到的是通用边缘纹理特征,学习率应该更小,越靠近输出层,学习率乘以的衰减系数越大。

4.2 optimizer.py 中 AdamW 与 Layer-wise decay 的耦合

源码包里的optimizer.py不像普通分类网络那样直接给整个模型设置同一份学习率,而是按 Transformer 的层索引对参数分组,每一组的学习率逐层放大。

def get_parameter_groups(model, weight_decay, layer_decay): parameter_group_names = {} parameter_group_vars = {} for name, param in model.named_parameters(): if not param.requires_grad: continue # 按层深度分配 decay 系数 if 'layers.' in name: depth = int(name.split('layers.')[1].split('.')[0]) this_layer_decay = layer_decay ** (num_layers - depth - 1) else: this_layer_decay = 1.0 # 区分 weight 和 bias,以及 norm 层 group_name = f"layer_{depth}_decay_{this_layer_decay}" if 'bias' in name or 'norm' in name: this_weight_decay = 0.0 else: this_weight_decay = weight_decay if group_name not in parameter_group_vars: parameter_group_vars[group_name] = { 'params': [], 'weight_decay': this_weight_decay, 'lr': base_lr * this_layer_decay, } parameter_group_vars[group_name]['params'].append(param) return list(parameter_group_vars.values())

这里最关键的是layer_decay ** (num_layers - depth - 1)这个指数映射:越靠近输入的层,depth 越小,幂次越大,但底数小于 1,所以学习率反而越小。用 Swin-T 的 4 个 stage 来算,如果layer_decay=0.65,第 3 个 stage 的层参数学习率大约是第 0 个 stage 的 0.65 倍,这个差距在训练后期会显著影响微调效果。

提示:设置weight_decay=0.05时,bias 和 norm 层会被排除在 weight decay 之外。这是 ViT 系列微调的通用实践,但初次实验如果发现 top-1 始终上不去,可以先检查一下 norm 层是否因为参数分组不对而产生了过大的梯度。

4.3 lr_scheduler.py 与 warmup 的交互

MoBY 的调度器用的是余弦退火,但前 20 个 epoch 内学习率会从 0 线性上升到base_lr。这部分代码不复杂,但和 MoCo v3 的报错日志对比看能发现一个问题:warmup 的作用不只是让训练曲线好看,而是避免大 batch size 下 AdamW 的方差估计在前期剧烈波动。如果直接去掉 warmup,top-1 会下降 0.5 到 1 个点。

# lr_scheduler.py 中的核心逻辑(简化) def adjust_learning_rate(optimizer, epoch, args): if epoch < args.warmup_epochs: # 线性 warmup lr = args.base_lr * (epoch + 1) / args.warmup_epochs else: # 余弦退火,T 是总 epoch progress = (epoch - args.warmup_epochs) / (args.max_epochs - args.warmup_epochs) lr = args.base_lr * 0.5 * (1.0 + math.cos(math.pi * progress)) for param_group in optimizer.param_groups: param_group['lr'] = lr

需要留意的是,optimizer.param_groups里的lr字段本来在optimizer.py里已经按 layer_decay 做了分组缩放,这里再整体乘以各层系数,相当于每一层的实际学习率是base_lr * layer_decay^depth * cosine_coeff。所以如果你在某个 layer 的学习率上手动做了修改,记得同步检查余弦退火的起点,避免重复缩放。

5. 线性评估与训练稳定性排查技巧

5.1 moby_linear.py 的评估流程

自监督学习的产出是特征提取器,判定特征质量的通用手段是线性评估:把模型主干冻结,只训练一个全新的全连接层。moby_linear.py里默认只保留 backbone 的输出,然后接一个nn.Linear(embed_dim, num_classes),用比预训练大得多的学习率去训。

# 评估的推荐命令(batch size 按单卡 256) python moby_linear.py \ --config configs/moby_swin_tiny.yaml \ --pretrained /path/to/checkpoint.pth \ --batch-size 256 \ --lr 30.0 \ --epochs 90

线性评估的lr是 30.0,这个数值在监督训练里看起来大得离谱,但对冻结特征提取器的线性层来说刚好合适。原因是 backbone 不参与反向传播,梯度方差小,用大学习率能更快逼近线性分类面的最优解。如果线性评估的 top-1 在 75 个 epoch 后还在上涨,说明预训练的表征还没充分拟合,可以适当延长评估 epoch。

5.2 指标对照:DeiT-S 与 Swin-T 的结果解读

主干网络预训练 epoch线性评估 top-1备注
DeiT-S30072.8%较轻量,适合单卡调试
Swin-T (w/o BN)30074.2%默认配置
Swin-T (w/ BN)30075.0%带 BatchNorm 更稳

这组数据在README.md里有直接对应。在 Swin-T 上启用 BatchNorm 后,线性评估的准确率提高了近 1 个点,但代价是 batch size 的敏感度上升。如果显存只允许每卡 64 的 batch,建议取消 BN,否则 BN 统计量会因为 batch 过小而失真。

5.3 训练崩到 loss 为 NaN 时的排查顺序

自监督 Transformer 训练最常见的崩法就是 loss 变 NaN,出现后按以下顺序检查。

先看学习率。如果base_lr没有按 batch size 做线性缩放,AdamW 很容易在训练中期爆梯度,建议把初始学习率缩半跑 10 个 epoch 做对照。再看队列尺寸。moby.py里 queue 默认是 65536,l_neg的计算量随队列长度线性增长,如果显存不够导致 batch size 被迫缩小,保底方案是把temperature从 0.2 提到 0.3,让 logits 的数值范围收窄,减缓梯度尖峰。

最后看数据增强。MoBY 第二个视图如果旋转角度太大,小物体上的语义信息会被破坏,模型可能学习到"靠着颜色统计来分辨",在验证集上表现不会崩,但训练 loss 会在 120 epoch 后开始抖动。对齐custom_image_folder.py里两个 transform 的参数,这是最后一步。

5.4 队列指针与 checkpoint 恢复的隐蔽问题

MoBY 把 queue 和queue_ptr都保存在了 checkpoint 里,但如果你的分布式训练中 rank 的卡不完全一致,恢复 checkpoint 时 queue 里的特征可能来自不同的数据分布,这就造成前几十个 batch 的负样本无意义。常见做法是恢复训练时把 queue 直接清空,让前 50 个 batch 作为预热。源码包里虽然没有显式做这个开关,但可以在moby_main.pyload_checkpoint逻辑里加一个reset_queue参数来规避。

另一个隐蔽点在zipreader.pycached_image_folder.py的配合上。如果开了缓存模式,会把每个 zip 内的读取结果缓存到内存,第一次 epoch 之后的速度会快很多,但缓存键是文件名加 zip 路径。如果你的训练脚本在同一个进程内切换了数据集路径,缓存不会自动失效,可能会导致 loader 一直返回旧数据。此时需要手动调用cache.clear()

掌握了这几个排查点,MoBY 的复现难度其实比 MoCo v3 要低不少,因为它把最小可行实现收敛在三个 yaml 和两个 Python 文件里,剩下的就是耐住性子调 BN 和 layer_decay 这两个参数。

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

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询