基于视觉Transformer的自监督学习

作者：Zhenda Xie*, Yutong Lin*, Zhuliang Yao, Zheng Zhang, Qi Dai, Yue Cao 和 Han Hu

本仓库是论文《基于Swin Transformer的自监督学习》（arXiv:2105.04553）的官方实现。

该代码库的一个重要特性是将Swin Transformer作为骨干网络之一，从而能够在目标检测和语义分割等下游任务上评估所学表征的迁移性能。 由于此前的工作大多使用ViT/DeiT，而这些模型在下游任务上的表现尚未得到充分优化，因此通常不会包含此类评估。

目前，本仓库包含了以下任务的相关代码和模型：

自监督学习与线性评估：已包含在本仓库中。请参阅 get_started.md 以快速入门。

目标检测/实例分割的迁移性能：请参阅 Swin Transformer用于目标检测。

语义分割的迁移性能：请参阅 Swin Transformer用于语义分割。

亮点

• 包含下游评估：首次针对基于Transformer的自监督学习，在下游任务上评估其迁移性能。
• 精简的技巧：相比MoCo v3和DINO等先前工作，使用的技巧显著减少。
• ImageNet-1K线性评估中的高精度：使用DeiT-S/16并在300个epoch上预训练时，准确率为72.8%，优于MoCo v3的72.5%和DINO的72.5%。

更新

2021年5月13日

1. 提供了在ImageNet-1K数据集上使用DeiT-Small进行自监督学习的模型（MoBY-DeiT-Small-300Ep-Pretrained，MoBY-DeiT-Small-300Ep-Linear)。
2. 提供了支持DeiT-Small自监督学习的代码和配置文件。

2021年5月11日

初始提交：

1. 提供了在ImageNet-1K数据集上进行自监督预训练的模型（MoBY-Swin-T-300Ep-Pretrained，MoBY-Swin-T-300Ep-Linear)。
2. 提供了支持自监督预训练、ImageNet-1K线性评估、COCO目标检测以及ADE20K语义分割的相关代码和模型。

简介

MoBY：结合MoCo v2和BYOL的自监督学习方法

MoBY（名称“MoBY”代表“Mo”Co v2与“BY”OL的结合）最初在arXiv论文（arXiv:2105.04553）中被提出，它是两种流行的自监督学习方法——MoCo v2和BYOL——的结合。它继承了MoCo v2中的动量设计、关键队列和对比损失，同时也借鉴了BYOL中的非对称编码器、非对称数据增强以及动量调度策略。

MoBY 在ImageNet-1K数据集上的线性评估中取得了相当高的准确率：分别使用DeiT和Swin-T作为骨干网络时，300个epoch的训练可达到72.8%和75.3%的Top-1准确率。这一性能与近期采用DeiT作为骨干网络的MoCo v3和DINO等方法相当，但所需的技巧却少得多。

Swin Transformer作为骨干网络

Swin Transformer（名称“Swin”意为“移位窗口”）最初在arXiv论文（arXiv:2103.14030）中被提出，它能够出色地作为计算机视觉领域的通用骨干网络。在COCO目标检测任务中，其表现尤为突出（test-dev上的box AP为58.7，mask AP为51.1）；而在ADE20K语义分割任务中，mIoU达到了53.5%，远超之前的模型。

我们选择Swin Transformer作为骨干网络之一，以评估其在目标检测等下游任务中的迁移性能。这使得本代码库区别于其他研究Transformer架构下自监督学习的方法。

ImageNet-1K 线性评估

方法	架构	Epochs	参数量	FLOPs	img/s	Top-1 准确率	预训练检查点	线性检查点
有监督	Swin-T	300	28M	4.5G	755.2	81.2	这里
MoBY	Swin-T	100	28M	4.5G	755.2	70.9	待定
MoBY1	Swin-T	100	28M	4.5G	755.2	72.0	待定
MoBY	DeiT-S	300	22M	4.6G	940.4	72.8	GoogleDrive/GitHub/百度网盘	GoogleDrive/GitHub/百度网盘
MoBY	Swin-T	300	28M	4.5G	755.2	75.3	GoogleDrive/GitHub/百度网盘	GoogleDrive/GitHub/百度网盘

• ¹ 表示采用了 MoCo v3 中的一个技巧，即用 BatchNorm 层替换 MLP 块之前的 LayerNorm 层的 MoBY 结果。

• 百度网盘的提取码为 moby。

迁移到下游任务

COCO 目标检测（2017 val）

主干网络	方法	模型	时序	box mAP	mask mAP	参数量	FLOPs
Swin-T	Mask R-CNN	有监督	1x	43.7	39.8	48M	267G
Swin-T	Mask R-CNN	MoBY	1x	43.6	39.6	48M	267G
Swin-T	Mask R-CNN	有监督	3x	46.0	41.6	48M	267G
Swin-T	Mask R-CNN	MoBY	3x	46.0	41.7	48M	267G
Swin-T	Cascade Mask R-CNN	有监督	1x	48.1	41.7	86M	745G
Swin-T	Cascade Mask R-CNN	MoBY	1x	48.1	41.5	86M	745G
Swin-T	Cascade Mask R-CNN	有监督	3x	50.4	43.7	86M	745G
Swin-T	Cascade Mask R-CNN	MoBY	3x	50.2	43.5	86M	745G

ADE20K 语义分割（val）

主干网络	方法	模型	裁剪尺寸	时序	mIoU	mIoU (ms+flip)	参数量	FLOPs
Swin-T	UPerNet	有监督	512x512	160K	44.51	45.81	60M	945G
Swin-T	UPerNet	MoBY	512x512	160K	44.06	45.58	60M	945G

引用 MoBY 和 Swin

MoBY

@article{xie2021moby,
  title={Self-Supervised Learning with Swin Transformers}, 
  author={Zhenda Xie and Yutong Lin and Zhuliang Yao and Zheng Zhang and Qi Dai and Yue Cao and Han Hu},
  journal={arXiv preprint arXiv:2105.04553},
  year={2021}
}

Swin Transformer

@article{liu2021Swin,
  title={Swin Transformer: Hierarchical Vision Transformer using Shifted Windows},
  author={Liu, Ze and Lin, Yutong and Cao, Yue and Hu, Han and Wei, Yixuan and Zhang, Zheng and Lin, Stephen and Guo, Baining},
  journal={arXiv preprint arXiv:2103.14030},
  year={2021}
}

开始使用

• 对于 使用 MoBY 和 Swin Transformer 进行自监督预训练及线性评估，请参阅 get_started.md 获取详细说明。
• 对于 目标检测/实例分割性能迁移，请参阅 Swin Transformer for Object Detection。
• 对于 语义分割性能迁移，请参阅 Swin Transformer for Semantic Segmentation。

Transformer-SSL (MoBY) 快速上手指南

本指南基于官方仓库 Transformer-SSL，帮助用户快速使用 MoBY 算法结合 Swin Transformer 或 DeiT 进行自监督学习预训练及线性评估。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04+)
• Python: 3.7 或更高版本
• GPU: 支持 CUDA 的 NVIDIA GPU (建议显存 16GB+)
• 核心依赖:
  • PyTorch >= 1.8.0
  • torchvision
  • timm
  • mmcv-full (用于下游任务评估，若仅做预训练可暂不安装)
  • Apex (可选，用于混合精度加速)

前置检查： 请确保已安装正确的 CUDA 版本和对应的 PyTorch。

安装步骤

1. 克隆仓库

git clone https://github.com/SwinTransformer/Transformer-SSL.git
cd Transformer-SSL

2. 创建虚拟环境并安装依赖

推荐使用 Conda 管理环境：

conda create -n moby python=3.8 -y
conda activate moby

# 安装 PyTorch (根据实际 CUDA 版本调整，此处以 CUDA 11.1 为例)
pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html

# 安装其他核心依赖
pip install timm==0.4.12 opencv-python-headless scipy pandas

# 【国内加速】安装 mmcv-full (若需进行下游任务评估)
# 使用 OpenMMLab 清华镜像源加速安装
pip install mmcv-full==1.3.17 -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.9.0/index.html

3. 编译自定义算子 (如需)

如果项目中包含自定义 CUDA 算子，通常需要编译：

python setup.py build_ext --inplace

(注：MoBY 主要依赖标准 PyTorch 模块，若直接运行预训练脚本报错缺失算子，请执行此步)

基本使用

本仓库核心功能为 自监督预训练 (Self-Supervised Pre-training) 和 线性评估 (Linear Evaluation)。以下以最常用的 DeiT-Small backbone 在 ImageNet-1K 上的训练为例。

1. 数据准备

请将 ImageNet-1K 数据集整理为以下目录结构：

/path/to/imagenet/
  train/
    class1/
      img1.jpeg
    class2/
      img2.jpeg
  val/
    class1/
      img3.jpeg
    class2/
      img4.jpeg

2. 自监督预训练 (Pre-training)

使用 MoBY 算法训练 DeiT-Small 模型（300 epochs）。

python main_pretrain.py \
  --data-path /path/to/imagenet \
  --model deit_small_patch16_224 \
  --batch-size 256 \
  --epochs 300 \
  --warmup-epochs 10 \
  --blr 1.5e-4 \
  --weight-decay 0.05 \
  --output_dir ./output/pretrain_deit_s

• --model: 可选择 deit_small_patch16_224 或 swin_tiny_patch4_window7_224。
• --blr: 基础学习率 (base learning rate)，公式通常为 batch_size / 256 * base_lr。

3. 线性评估 (Linear Evaluation)

使用预训练好的权重冻结骨干网络，仅训练分类头进行线性探测。

python main_linear.py \
  --data-path /path/to/imagenet \
  --model deit_small_patch16_224 \
  --batch-size 256 \
  --epochs 90 \
  --blr 0.1 \
  --pretrained_weights ./output/pretrain_deit_s/checkpoint.pth \
  --output_dir ./output/linear_deit_s

4. 加载官方预训练模型

如果您不想从头训练，可以直接下载官方提供的预训练权重进行线性评估或下游任务迁移。

下载链接 (含国内百度盘加速):

• DeiT-S (300 Ep) 预训练权重: GoogleDrive | BaiduPan (提取码: moby)
• Swin-T (300 Ep) 预训练权重: GoogleDrive | BaiduPan (提取码: moby)

使用示例: 将下载的 .pth 文件路径填入上述线性评估命令的 --pretrained_weights 参数即可。

---

更多详细配置（如 Swin Transformer 的具体参数、多节点训练等）请参考仓库根目录下的 get_started.md 文件或配置文件目录 configs/。