在 MMDetection 中使用自监督预训练骨干:MoCo v2 与 SwAV 权重的转换、配置与训练实践
【免费下载链接】mmdetectionOpenMMLab Detection Toolbox and Benchmark项目地址: https://gitcode.com/gh_mirrors/mm/mmdetection
本文基于 MMDetection 仓库中configs/selfsup_pretrain/目录的官方文档与实践配置,系统讲解如何将 MoCo v2、SwAV 等自监督算法预训练出的 ResNet-50 骨干接入 Mask R-CNN 检测框架。读者将掌握从自监督权重下载、键名转换、骨干配置修改到完整训练评估的全流程,并理解frozen_stages、SyncBN、norm_eval等关键参数在源码层面的真实作用。
一、背景:为什么使用自监督预训练骨干
无监督图像表示学习已经显著缩小了与监督式预训练之间的差距,尤其是近期对比学习方法(contrastive learning)取得的进展。本目录聚焦两种代表性自监督算法产出的骨干权重:
- SwAV(Swapping Assignments between Views):一种在线算法,无需像传统对比学习那样计算大量显式的成对特征比较。其核心思想是:同时进行数据聚类,并强制同一图像不同增强视图(views)产生的聚类分配保持一致;具体实现为“交换预测机制”——从一个视图的表示去预测另一个视图的聚类分配。SwAV 支持大批量与小批量训练,可扩展到无限数据量;相比早期对比方法,它不需要大容量 memory bank,也不需要专门的动量网络,内存效率更高。同时它提出 multi-crop 数据增强策略,用不同分辨率的视图混合替代两个全分辨率视图,几乎不增加内存与算力开销。论文在 ImageNet 上用 ResNet-50 取得 75.3% top-1 精度,并在所有考察的下游迁移任务上超过监督预训练。
- MoCo(Momentum Contrast):将对比学习视为“字典查找”问题,通过维护一个队列(queue)和动量更新的编码器(moving-averaged encoder)构建动态、一致且足够大的字典,从而高效地支持对比式无监督学习。MoCo 在 ImageNet 线性评估协议下取得有竞争力的结果,更重要的是其学到的表示能很好地迁移到下游任务——在 PASCAL VOC、COCO 等数据集上的 7 个检测/分割任务中,MoCo 预训练均可超过其监督预训练对应版本,有时甚至以较大幅度胜出,说明在很多视觉任务中无监督与监督表示学习的差距已被大幅拉近。
MMDetection 通过 configs/selfsup_pretrain/README.md 提供了将上述自监督骨干应用于检测系统的官方支持,并在 Mask R-CNN 上给出完整结果。
二、总体使用流程:两步接入
在 MMDetection 中使用自监督预训练骨干,只需要两个步骤:
- 下载并转换模型:把 MoCo / SwAV 发布的权重转换成 MMDetection 所支持的 PyTorch 风格 checkpoint;
- 修改配置文件:替换骨干的初始化方式,并按自监督骨干的特性调整训练设置(主要是
frozen_stages与归一化层配置)。
下面分别展开。
三、第一步:模型权重转换(selfsup2mmdet.py)
3.1 转换脚本的通用用法
为保证通用性,MMDetection 在 tools 目录下提供了脚本selfsup2mmdet.py,用于将不同自监督方法预训练的模型键名转换为 MMDetection 使用的 PyTorch 风格 checkpoint。其用法为:
python -u tools/model_converters/selfsup2mmdet.py ${PRETRAIN_PATH} ${STORE_PATH} --selfsup ${method}${PRETRAIN_PATH}:自监督方法官方发布的原始权重路径;${STORE_PATH}:转换后 checkpoint 的保存路径;--selfsup ${method}:指定自监督方法,当前支持moco与swav两个选项。
3.2 MoCo v2 权重转换示例
要使用 MoCo 官方发布的 ResNet-50 骨干(MoCo v2,800 epochs 预训练),先下载其moco_v2_800ep_pretrain.pth.tar权重文件,然后执行:
python -u tools/model_converters/selfsup2mmdet.py ./moco_v2_800ep_pretrain.pth.tar mocov2_r50_800ep_pretrain.pth --selfsup moco命令执行后,会在当前目录生成mocov2_r50_800ep_pretrain.pth,这正是后续配置中init_cfg.checkpoint要引用的文件。
3.3 SwAV 权重转换示例
要使用 SwAV 发布的 ResNet-50 骨干,下载其swav_800ep_pretrain.pth.tar权重文件即可。从源码看,SwAV 的权重键名本身已经是 PyTorch 风格:
elif args.selfsup == 'swav': print('SWAV does not need to convert the keys')因此--selfsup swav分支不会改写文件,配置中可以直接把swav_800ep_pretrain.pth.tar原文件作为 checkpoint 使用(见 tools/model_converters/selfsup2mmdet.py)。
3.4 转换脚本的底层逻辑(源码解读)
从 tools/model_converters/selfsup2mmdet.py 的实现可以看出,MoCo 权重转换的核心逻辑位于moco_convert(src, dst)函数:
- 用
torch.load加载 MoCo 官方发布的权重文件,取出其中的state_dict; - 遍历全部键,只保留以
module.encoder_q.为前缀的条目——这正是 MoCo 动量编码器中 query 编码器(用于下游迁移的骨干)部分; - 将前缀
module.encoder_q.去除,得到形如conv1.weight、layer1.0.conv1.weight等标准 PyTorch ResNet 键名; - 将过滤后的
state_dict重新打包为{'state_dict': state_dict}的 checkpoint 并保存。
这一步骤本质上完成了“MoCo 训练时封装层(module.encoder_q.)与 MMDetection 骨干模块名(裸 ResNet 键名)”之间的命名空间对齐,转换后的权重即可被 configs/base/models/mask-rcnn_r50_fpn.py 中的 ResNet 骨干直接加载。
四、第二步:修改配置文件接入自监督骨干
4.1 关键差异:与默认 Mask R-CNN 骨干配置的对比
默认的 Mask R-CNN ResNet-50 FPN 骨干配置(见 configs/base/models/mask-rcnn_r50_fpn.py)为:
backbone=dict( type='ResNet', depth=50, num_stages=4, out_indices=(0, 1, 2, 3), frozen_stages=1, norm_cfg=dict(type='BN', requires_grad=True), norm_eval=True, style='pytorch', init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet50')),而自监督预训练骨干要求的三处关键修改是:
frozen_stages从1改为0:自监督预训练覆盖了包括 stem 在内的全部层,因此不再冻结浅层,让全部骨干参数参与微调;norm_cfg从type='BN'改为type='SyncBN':MoCo / SwAV 预训练时骨干使用 SyncBN(同步批归一化),微调阶段保持一致的归一化层类型才能正确加载统计量并维持稳定的训练行为;norm_eval从True改为False:让归一化层的 running statistics 在训练中继续更新(自监督预训练骨干的 BN 统计量也需要在检测数据上重新适应),同时保持requires_grad=True使其可被梯度更新。
4.2 使用 MoCo 骨干的完整配置示例
以下配置使用 MoCo 预训练骨干训练 Mask R-CNN(1x 单尺度,12 epochs):
_base_ = [ '../_base_/models/mask-rcnn_r50_fpn.py', '../_base_/datasets/coco_instance.py', '../_base_/schedules/schedule_1x.py', '../_base_/default_runtime.py' ] model = dict( pretrained='./mocov2_r50_800ep_pretrain.pth', backbone=dict( frozen_stages=0, norm_cfg=dict(type='SyncBN', requires_grad=True), norm_eval=False))这是 configs/selfsup_pretrain/README.md 中给出的经典写法(顶层pretrained字段会作为全局预训练权重加载)。
4.3 仓库中的四份官方配置
configs/selfsup_pretrain/目录实际提供四份可直接使用的配置,均采用init_cfg(3.x 推荐的初始化方式)而不是顶层pretrained字段:
| 配置 | 骨干来源 | 训练计划 | 数据增强 |
|---|---|---|---|
| mask-rcnn_r50-mocov2-pre_fpn_1x_coco.py | MoCo v2 | 1x(12 epochs) | 默认单尺度 (1333, 800) |
| mask-rcnn_r50-mocov2-pre_fpn_ms-2x_coco.py | MoCo v2 | 2x(24 epochs) | 多尺度 [1333, 640] ~ [1333, 800] |
| mask-rcnn_r50-swav-pre_fpn_1x_coco.py | SwAV | 1x(12 epochs) | 默认单尺度 (1333, 800) |
| mask-rcnn_r50-swav-pre_fpn_ms-2x_coco.py | SwAV | 2x(24 epochs) | 多尺度 [1333, 640] ~ [1333, 800] |
以 MoCo 2x 配置为例,其完整内容为:
_base_ = [ '../_base_/models/mask-rcnn_r50_fpn.py', '../_base_/datasets/coco_instance.py', '../_base_/schedules/schedule_2x.py', '../_base_/default_runtime.py' ] model = dict( backbone=dict( frozen_stages=0, norm_cfg=dict(type='SyncBN', requires_grad=True), norm_eval=False, init_cfg=dict( type='Pretrained', checkpoint='./mocov2_r50_800ep_pretrain.pth'))) train_pipeline = [ dict(type='LoadImageFromFile', backend_args={{_base_.backend_args}}), dict(type='LoadAnnotations', with_bbox=True, with_mask=True), dict( type='RandomResize', scale=[(1333, 640), (1333, 800)], keep_ratio=True), dict(type='RandomFlip', prob=0.5), dict(type='PackDetInputs') ] train_dataloader = dict(dataset=dict(pipeline=train_pipeline))这里的train_pipeline覆盖了 configs/base/datasets/coco_instance.py 中的默认 pipeline,将固定尺度Resize(1333, 800)替换为在短边 640~800 之间随机采样的RandomResize,实现多尺度训练;其余数据加载、标注读取(with_bbox=True, with_mask=True,对应实例分割任务)逻辑保持不变。SwAV 的配置与之同构,仅checkpoint指向./swav_800ep_pretrain.pth.tar。
4.4 训练计划说明
两份 2x 配置基于 configs/base/schedules/schedule_2x.py:训练 24 个 epoch,MultiStepLR在 epoch 16 和 22 处衰减 10 倍,前 500 步使用LinearLR热身(start_factor=0.001),优化器为 SGD(lr=0.02、momentum=0.9、weight_decay=0.0001),并预设base_batch_size=16(8 GPU × 2 samples)作为自动缩放学习率的基准。1x 配置则基于 configs/base/schedules/schedule_1x.py,训练 12 个 epoch,衰减点在 epoch 8 和 11。
配置修改完成后,即可通过标准训练入口启动训练(单卡或分布式均可):
python tools/train.py configs/selfsup_pretrain/mask-rcnn_r50-mocov2-pre_fpn_1x_coco.py五、源码级理解:frozen_stages 与 norm_eval 到底做了什么
上述三处参数修改并非文档凭空规定,而是与骨干实现严格对应。查看 mmdet/models/backbones/resnet.py 的 ResNet 实现:
frozen_stages(默认-1,意为不冻结)控制冻结的骨干 stage 数量。在_freeze_stages()方法(mmdet/models/backbones/resnet.py)中:当frozen_stages >= 0时,stem(conv1+norm1,或deep_stem)被置为 eval 模式且参数requires_grad=False;随后for i in range(1, self.frozen_stages + 1)会把对应layer{i}同样冻结。默认配置frozen_stages=1意味着 stem 冻结、layer1冻结;而自监督场景改为0后,仅 stem 进入 eval 且停止梯度,layer1及以上全部参与微调。norm_eval(默认True)作用于train()方法(mmdet/models/backbones/resnet.py):当模型处于训练模式且norm_eval=True时,遍历所有子模块,将所有_BatchNorm实例置为 eval 模式(即不再更新 running mean/var,仅使用缓存统计量做前向)。这正是 ImageNet 监督预训练迁移到检测任务时的经典做法。而自监督骨干将其设为False,配合norm_cfg=dict(type='SyncBN', requires_grad=True),使 BN 统计量与参数在检测训练中持续更新。
从源码结构可以推断:之所以要求自监督骨干必须配SyncBN,是因为 MoCo / SwAV 预训练阶段即采用同步批归一化,若微调时改用普通BN,归一化层的语义(跨卡统计)发生变化,会影响预训练知识的高效迁移。
六、实验结果:Mask R-CNN 上的自监督预训练表现
configs/selfsup_pretrain/README.md 给出了在 COCO 实例分割任务上的官方结果(测试环境为标准 COCO 评估协议):
| Method | Backbone | Style | Lr schd | box AP | mask AP | Config |
|---|---|---|---|---|---|---|
| Mask RCNN | R50 by MoCo v2 | pytorch | 1x | 38.0 | 34.3 | config |
| Mask RCNN | R50 by MoCo v2 | pytorch | multi-scale 2x | 40.8 | 36.8 | config |
| Mask RCNN | R50 by SwAV | pytorch | 1x | 39.1 | 35.7 | config |
| Mask RCNN | R50 by SwAV | pytorch | multi-scale 2x | 41.3 | 37.3 | config |
对应各配置的预训练模型权重与训练日志文件均可在 OpenMMLab 官方模型库对应模型页面下载,文件名形如mask_rcnn_r50_fpn_mocov2-pretrain_1x_coco_20210604_114614-a8b63483.pth与同名.log.json。
重要说明(来自官方文档 Notice):以上仅提供单尺度 1x 与多尺度 2x 配置作为“如何接入自监督骨干”的示例;官方后续计划使用发布的自监督骨干复现对应论文中的完整结果,读者可以持续关注仓库更新。
七、引用与致谢
MMDetection 支持将不同自监督方法预训练的骨干应用到检测系统中,并在 Mask R-CNN 上提供结果。预训练模型由 MoCo 转换而来、并从 SwAV 官方发布下载。若在研究中使用了这些工作,请引用:
SwAV:
@article{caron2020unsupervised, title={Unsupervised Learning of Visual Features by Contrasting Cluster Assignments}, author={Caron, Mathilde and Misra, Ishan and Mairal, Julien and Goyal, Priya and Bojanowski, Piotr and Joulin, Armand}, booktitle={Proceedings of Advances in Neural Information Processing Systems (NeurIPS)}, year={2020} }MoCo 与 MoCo v2:
@Article{he2019moco, author = {Kaiming He and Haoqi Fan and Yuxin Wu and Saining Xie and Ross Girshick}, title = {Momentum Contrast for Unsupervised Visual Representation Learning}, journal = {arXiv preprint arXiv:1911.05722}, year = {2019}, } @Article{chen2020mocov2, author = {Xinlei Chen and Haoqi Fan and Ross Girshick and Kaiming He}, title = {Improved Baselines with Momentum Contrastive Learning}, journal = {arXiv preprint arXiv:2003.04297}, year = {2020}, }八、小结与扩展阅读
接入自监督预训练骨干只需“转换权重 + 修改配置”两步,核心要点可总结为:
- 用 tools/model_converters/selfsup2mmdet.py 将 MoCo 权重的
module.encoder_q.前缀去除(SwAV 无需转换); - 在骨干配置中设置
frozen_stages=0、norm_cfg=dict(type='SyncBN', requires_grad=True)、norm_eval=False,并通过init_cfg=dict(type='Pretrained', checkpoint=...)加载转换后的权重; - 多尺度训练时按需覆盖
train_pipeline的RandomResize,并搭配 2x 训练计划以获得更高精度。
本目录只是起点——同样的接入模式可以推广到任意自监督算法产出的骨干。若希望深入理解配置继承机制,可继续阅读 configs/base/models/mask-rcnn_r50_fpn.py、configs/base/datasets/coco_instance.py 与 configs/base/schedules/schedule_1x.py 三份基础配置;若想掌握骨干实现的细节,可研读 mmdet/models/backbones/resnet.py 中_freeze_stages与train两个方法的完整实现。
【免费下载链接】mmdetectionOpenMMLab Detection Toolbox and Benchmark项目地址: https://gitcode.com/gh_mirrors/mm/mmdetection
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考