solo-learn

一个基于 PyTorch Lightning 的自监督方法库，用于无监督视觉表征学习。我们的目标是在可比较的环境中提供当前最优的自监督方法，并同时实现一系列训练技巧。该库是自包含的，但用户也可以在 solo-learn 外部使用其中的模型。更多详情请参阅我们的论文。

---

最新动态

• [2024年1月14日]：:clap: 2023年期间进行了大量稳定性改进 :) 同时新增了 All4One。
• [2023年1月7日]：:diving_mask: 添加了 MAE 在 ImageNet 上的结果、检查点和配置文件。感谢 HuangChiEn 的贡献。
• [2022年12月31日]：:stars: 全新Logo！特别感谢 Luiz！
• [2022年9月27日]：:pencil: 全新的配置系统，采用 OmegaConf/Hydra 框架，提升了清晰度和灵活性。后续将推出更多教程！
• [2022年8月4日]：:paintbrush: 新增 MAE，并支持通过 main_linear.py 对主干网络进行微调，以及 mixup、cutmix 和 随机增强。
• [2022年7月13日]：:sparkling_heart: 增加了对 H5 数据的支持，优化了脚本和数据处理流程。
• [2022年6月26日]：:fire: 新增了 MoCo V3。
• [2022年6月10日]：:bomb: 改进了 LARS 优化器。
• [2022年6月9日]：:lollipop: 增加了对 WideResnet 的支持，为 SwAV 添加了多尺度裁剪功能，并引入了均衡化数据增强技术。
• [2022年5月2日]：:diamond_shape_with_a_dot_inside: 使用 DataModule 封装了 Dali，增加了线性评估的自动恢复功能以及 Wandb 运行的恢复功能。
• [2022年4月12日]：:rainbow: 改进了模型设计，并支持仅使用部分数据进行训练。
• [2022年4月1日]：:mag: 新增了使用 通道最后转换 的选项，这能显著缩短训练时间。
• [2022年2月4日]：:partying_face: 论文已被 JMLR 接受。
• [2022年1月31日]：:eye: 增加了与 timm 集成的 ConvNeXt 支持。
• [2021年12月20日]：:thermometer: 添加了 MoCo V2+ 在 ImageNet 上的结果、脚本和检查点。
• [2021年12月5日]：:notes: 将 SupCon 从 SimCLR 中分离出来，并添加了相关运行记录。
• [2021年12月1日]：:fountain: 新增了 PoolFormer。
• [2021年11月29日]：:bangbang: 重大变更！请更新版本！！！
• [2021年11月29日]：:book: 新增教程！
• [2021年11月29日]：:houses: 新增了离线 K-NN 和离线 UMAP 功能。
• [2021年11月29日]：:rotating_light: 更新了 PyTorch 和 PyTorch Lightning 的版本，速度提升了约 10%。
• [2021年11月29日]：:beers: 添加了行为准则、贡献指南、问题模板以及 UMAP 教程。
• [2021年11月23日]：:space_invader: 新增了 VIbCReg。
• [2021年10月21日]：:triumph: 增加了通过 Detectron v2 进行目标检测的支持，并实现了自动恢复功能，能够自动尝试恢复因崩溃或超时而中断的实验。
• [2021年10月10日]：:japanese_ogre: 重构了数据增强流水线，以提高灵活性并支持多尺度裁剪。同时为 BYOL 添加了多尺度裁剪功能。
• [2021年9月27日]：:pizza: 新增了 NNSiam、NNBYOL，并推出了关于如何实现新方法的新教程 1 和 2，进一步加强了测试，并修复了自定义数据和线性评估方面的问题。
• [2021年9月19日]：:kangaroo: 新增了在线 k-NN 评估功能。
• [2021年9月17日]：:robot: 新增了 ViT 和 Swin。
• [2021年9月13日]：:book: 改进了 文档，并新增了关于 预训练 和 离线线性评估 的教程。
• [2021年8月13日]：:whale: DeepCluster V2 现已可用。

---

路线图及所需帮助

• 重新整理文档以提升清晰度。
• 提供更好、更及时的教程。
• 增加性能相关的测试，确保方法在每次更新后都能保持一致的性能。
• 持续添加新方法。

---

可用方法

• All4One
• Barlow Twins
• BYOL
• DeepCluster V2
• DINO
• MAE
• MoCo V2+
• MoCo V3
• NNBYOL
• NNCLR
• NNSiam
• ReSSL
• SimCLR
• SimSiam
• 有监督对比学习
• SwAV
• VIbCReg
• VICReg
• W-MSE

---

额外特色

主干网络

• ResNet
• WideResNet
• ViT
• Swin
• PoolFormer
• ConvNeXt

数据处理

• 使用 Nvidia Dali 将数据处理速度提升高达 100%。
• 灵活的数据增强技术。

评估

• 通过停止梯度进行在线线性评估，便于调试和原型开发（也可选择在动量骨干网络上使用）。
• 标准的离线线性评估。
• 在线和离线 K-NN 评估。
• 使用 UMAP 自动可视化特征空间。

训练技巧

• PyTorch Lightning 的所有优势（混合精度、梯度累积、梯度裁剪等）。
• 通道最后转换。
• 按照 SwAV 实现的多裁剪数据加载：
  • 注意：目前仅 SimCLR、BYOL 和 SwAV 支持此功能。
• 将批归一化层和偏置项排除在权重衰减和 LARS 优化之外。
• 投影头不使用学习率调度器（如 SimSiam 中的做法）。

日志记录

• 使用 WandB 在云端记录指标。
• 自定义模型检查点保存，并采用简洁的文件组织方式。

---

需求

• torch
• torchvision
• tqdm
• einops
• wandb
• pytorch-lightning
• lightning-bolts
• torchmetrics
• scipy
• timm

可选：

• nvidia-dali
• matplotlib
• seaborn
• pandas
• umap-learn

---

安装

首先克隆仓库。

然后，若需安装支持 Dali 和/或 UMAP 的 solo-learn，可运行以下命令：

pip3 install .[dali,umap,h5] --extra-index-url https://developer.download.nvidia.com/compute/redist

若无需 Dali/UMAP/H5 支持，则直接安装仓库即可：

pip3 install .

本地开发时：

pip3 install -e .[umap,h5]
# 确保已安装 pre-commit 钩子
pre-commit install

注意：若遇到 Dali 安装问题，请按照其官方 指南 进行安装。

注意 2：建议安装 Pillow-SIMD，以便在不使用 Dali 时提升加载速度。

注意 3：即将上线 PyPI。

---

训练

预训练骨干网络时，可参考 scripts/pretrain/ 目录下的多个 Bash 脚本。我们目前使用 Hydra 处理配置文件，常见用法如下：

python3 main_pretrain.py \
    # 指定训练脚本所在目录
    --config-path scripts/pretrain/imagenet-100/ \
    # 指定训练配置文件名
    --config-name barlow.yaml
    # 可通过 ++new_argument=VALUE 添加新参数（例如未在 YAML 文件中定义的参数）
    # 也可在此处添加 PyTorch Lightning 的相关参数。

完成预训练后，可参照 scripts/linear 或 scripts/finetune 中的示例进行离线线性评估或对整个骨干网络进行微调。

K-NN 评估和 UMAP 可视化则可在 scripts/{knn,umap} 目录下的脚本中找到。

注意：这些脚本会尽量保持最新，并尽可能遵循各论文推荐的超参数，但请在运行前仔细核对。

---

教程

请查看我们的 文档 和教程：

• 概述
• 离线线性评估
• 目标检测
• 添加新方法
• 添加动量方法
• 使用 UMAP 可视化特征
• 离线 K-NN

若希望为 solo-learn 做贡献，请务必阅读 如何贡献 并遵守 行为准则。

---

模型库

所有预训练模型均可通过下方表格直接下载，或通过运行以下脚本以编程方式获取： zoo/cifar10.sh、zoo/cifar100.sh、zoo/imagenet100.sh 和 zoo/imagenet.sh。

---

结果

注意：部分超参数可能并非最优，我们将在后续重新运行表现较差的方法。

CIFAR-10

方法	主干网络	Epochs	Dali	Acc@1	Acc@5	检查点
All4One	ResNet18	1000	:x:	93.24	99.88	:link:
Barlow Twins	ResNet18	1000	:x:	92.10	99.73	:link:
BYOL	ResNet18	1000	:x:	92.58	99.79	:link:
DeepCluster V2	ResNet18	1000	:x:	88.85	99.58	:link:
DINO	ResNet18	1000	:x:	89.52	99.71	:link:
MoCo V2+	ResNet18	1000	:x:	92.94	99.79	:link:
MoCo V3	ResNet18	1000	:x:	93.10	99.80	:link:
NNCLR	ResNet18	1000	:x:	91.88	99.78	:link:
ReSSL	ResNet18	1000	:x:	90.63	99.62	:link:
SimCLR	ResNet18	1000	:x:	90.74	99.75	:link:
Simsiam	ResNet18	1000	:x:	90.51	99.72	:link:
SupCon	ResNet18	1000	:x:	93.82	99.65	:link:
SwAV	ResNet18	1000	:x:	89.17	99.68	:link:
VIbCReg	ResNet18	1000	:x:	91.18	99.74	:link:
VICReg	ResNet18	1000	:x:	92.07	99.74	:link:
W-MSE	ResNet18	1000	:x:	88.67	99.68	:link:

CIFAR-100

方法	主干网络	Epochs	Dali	Acc@1	Acc@5	检查点
All4One	ResNet18	1000	:x:	72.17	93.35	:link:
Barlow Twins	ResNet18	1000	:x:	70.90	91.91	:link:
BYOL	ResNet18	1000	:x:	70.46	91.96	:link:
DeepCluster V2	ResNet18	1000	:x:	63.61	88.09	:link:
DINO	ResNet18	1000	:x:	66.76	90.34	:link:
MoCo V2+	ResNet18	1000	:x:	69.89	91.65	:link:
MoCo V3	ResNet18	1000	:x:	68.83	90.57	:link:
NNCLR	ResNet18	1000	:x:	69.62	91.52	:link:
ReSSL	ResNet18	1000	:x:	65.92	89.73	:link:
SimCLR	ResNet18	1000	:x:	65.78	89.04	:link:
Simsiam	ResNet18	1000	:x:	66.04	89.62	:link:
SupCon	ResNet18	1000	:x:	70.38	89.57	:link:
SwAV	ResNet18	1000	:x:	64.88	88.78	:link:
VIbCReg	ResNet18	1000	:x:	67.37	90.07	:link:
VICReg	ResNet18	1000	:x:	68.54	90.83	:link:
W-MSE	ResNet18	1000	:x:	61.33	87.26	:link:

ImageNet-100

方法	主干网络	轮数	Dali	精确率@1（在线）	精确率@1（离线）	精确率@5（在线）	精确率@5（离线）	检查点
All4One	ResNet18	400	:heavy_check_mark:	81.93	-	96.23	-	:link:
Barlow Twins :rocket:	ResNet18	400	:heavy_check_mark:	80.38	80.16	95.28	95.14	:link:
BYOL :rocket:	ResNet18	400	:heavy_check_mark:	80.16	80.32	95.02	94.94	:link:
DeepCluster V2	ResNet18	400	:x:	75.36	75.4	93.22	93.10	:link:
DINO	ResNet18	400	:heavy_check_mark:	74.84	74.92	92.92	92.78	:link:
DINO :sleepy:	ViT Tiny	400	:x:	63.04	TODO	87.72	TODO	:link:
MoCo V2+ :rocket:	ResNet18	400	:heavy_check_mark:	78.20	79.28	95.50	95.18	:link:
MoCo V3 :rocket:	ResNet18	400	:heavy_check_mark:	80.36	80.36	95.18	94.96	:link:
MoCo V3 :rocket:	ResNet50	400	:heavy_check_mark:	85.48	84.58	96.82	96.70	:link:
NNCLR :rocket:	ResNet18	400	:heavy_check_mark:	79.80	80.16	95.28	95.30	:link:
ReSSL	ResNet18	400	:heavy_check_mark:	76.92	78.48	94.20	94.24	:link:
SimCLR :rocket:	ResNet18	400	:heavy_check_mark:	77.64	TODO	94.06	TODO	:link:
Simsiam	ResNet18	400	:heavy_check_mark:	74.54	78.72	93.16	94.78	:link:
SupCon	ResNet18	400	:heavy_check_mark:	84.40	TODO	95.72	TODO	:link:
SwAV	ResNet18	400	:heavy_check_mark:	74.04	74.28	92.70	92.84	:link:
VIbCReg	ResNet18	400	:heavy_check_mark:	79.86	79.38	94.98	94.60	:link:
VICReg :rocket:	ResNet18	400	:heavy_check_mark:	79.22	79.40	95.06	95.02	:link:
W-MSE	ResNet18	400	:heavy_check_mark:	67.60	69.06	90.94	91.22	:link:

:rocket: 表示超参数经过大量调优的方法。

:sleepy: ViT 非常消耗计算资源且不稳定，因此我们正在逐步尝试更大规模的架构和更大的批量大小。目前总批量大小为 128，且需要使用 float32 精度。如果您想参与运行，请联系我们！

ImageNet

方法	主干网络	轮数	Dali	精确率@1（在线）	精确率@1（离线）	精确率@5（在线）	精确率@5（离线）	检查点	微调后的检查点
Barlow Twins	ResNet50	100	:heavy_check_mark:	67.18	67.23	87.69	87.98	:link:
BYOL	ResNet50	100	:heavy_check_mark:	68.63	68.37	88.80	88.66	:link:
MoCo V2+	ResNet50	100	:heavy_check_mark:	62.61	66.84	85.40	87.60	:link:
MAE	ViT-B/16	100	:x:	~	81.60（微调后）	~	95.50（微调后）	:link:	:link:?

DALI 的训练效率

我们在一台配备 Intel i9-9820X 处理器和两块 RTX 2080 Ti 显卡的服务器上，使用带有和不带 DALI 的 ResNet18 模型（每张 GPU 配置 4 个工作进程），报告了部分方法的训练效率。

方法	Dali	20 个 epoch 的总耗时	每个 epoch 的耗时	GPU 显存（每张 GPU）
Barlow Twins	:x:	1 小时 38 分 27 秒	4 分 55 秒	5097 MB
	:heavy_check_mark:	43 分 2 秒	2 分 10 秒（快 56%）	9292 MB
BYOL	:x:	1 小时 38 分 46 秒	4 分 56 秒	5409 MB
	:heavy_check_mark:	50 分 33 秒	2 分 31 秒（快 49%）	9521 MB
NNCLR	:x:	1 小时 38 分 30 秒	4 分 55 秒	5060 MB
	:heavy_check_mark:	42 分 3 秒	2 分 6 秒（快 64%）	9244 MB

注：GPU 显存占用的增加并不与模型大小成正比，而是与工作进程的数量相关。

---

引用

如果您使用 solo-learn，请引用我们的论文：

@article{JMLR:v23:21-1155,
  author  = {Victor Guilherme Turrisi da Costa 和 Enrico Fini 和 Moin Nabi 和 Nicu Sebe 和 Elisa Ricci},
  title   = {solo-learn：用于视觉表征学习的自监督方法库},
  journal = {机器学习研究期刊},
  year    = {2022},
  volume  = {23},
  number  = {56},
  pages   = {1-6},
  url     = {http://jmlr.org/papers/v23/21-1155.html}
}

solo-learn 快速上手指南

solo-learn 是一个基于 PyTorch Lightning 的自监督视觉表示学习库，集成了多种 SOTA（最先进）算法（如 SimCLR, BYOL, DINO, MAE 等），旨在提供可复现的训练环境和灵活的配置系统。

环境准备

系统要求

• 操作系统: Linux (推荐), macOS, Windows
• Python: 3.8+
• GPU: 支持 CUDA 的 NVIDIA GPU（可选，但强烈推荐用于训练）

前置依赖

核心依赖包括：

• torch & torchvision
• pytorch-lightning
• hydra-core (用于配置管理)
• wandb (可选，用于日志记录)

可选增强依赖（用于加速数据加载或可视化）：

• nvidia-dali: 大幅提升数据预处理速度
• umap-learn, matplotlib, seaborn: 用于特征空间可视化

国内开发者提示：建议配置 pip 国内镜像源（如清华源、阿里源）以加速依赖下载。

pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆仓库

git clone https://github.com/vturrisi/solo-learn.git
cd solo-learn

2. 安装依赖

方案 A：完整安装（推荐） 包含 Dali (数据加速)、UMAP (可视化) 和 H5 支持。 注意：安装 Dali 需要添加 NVIDIA 的额外索引源。

pip3 install .[dali,umap,h5] --extra-index-url https://developer.download.nvidia.com/compute/redist

方案 B：最小化安装 仅安装核心功能，不包含 Dali 和 UMAP。

pip3 install .

方案 C：开发模式安装 如果您计划修改源码或贡献代码。

pip3 install -e .[umap,h5]
pre-commit install

性能优化提示：如果不使用 Dali，建议安装 Pillow-SIMD 以提升图像加载速度：

pip uninstall pillow
CC="cc -mavx2" pip install -U --force-reinstall pillow-simd

基本使用

solo-learn 使用 Hydra 进行配置管理，通过命令行参数覆盖 YAML 配置文件中的设置。

1. 预训练骨干网络 (Pretraining)

进入项目根目录，使用 main_pretrain.py 启动训练。以下示例展示如何使用配置文件启动 Barlow Twins 算法在 ImageNet-100 数据集上的训练：

python3 main_pretrain.py \
    --config-path scripts/pretrain/imagenet-100/ \
    --config-name barlow.yaml \
    ++data.data_path=/path/to/your/dataset \
    ++trainer.gpus=1 \
    ++trainer.max_epochs=100

参数说明：

• --config-path: 配置文件所在文件夹路径（位于 scripts/pretrain/ 下）。
• --config-name: 具体的算法配置文件名（如 barlow.yaml, simclr.yaml, dino.yaml 等）。
• ++key=value: 动态覆盖配置项。例如指定数据路径、GPU 数量或 epoch 数。也可以直接传递 PyTorch Lightning 的参数。

常用算法配置文件位于 scripts/pretrain/ 目录下，支持 CIFAR-10, CIFAR-100, ImageNet-100 及完整 ImageNet。

2. 线性评估 (Linear Evaluation)

预训练完成后，使用冻结的骨干网络进行线性分类器评估：

python3 main_linear.py \
    --config-path scripts/linear/imagenet-100/ \
    --config-name barlow.yaml \
    ++model.checkpoint_path=/path/to/pretrained_checkpoint.ckpt \
    ++data.data_path=/path/to/your/dataset

3. 微调 (Finetuning)

如果需要微调整个骨干网络（例如使用 MAE 或其他需要微调的方法）：

python3 main_linear.py \
    --config-path scripts/finetune/imagenet-100/ \
    --config-name mae.yaml \
    ++model.checkpoint_path=/path/to/pretrained_checkpoint.ckpt

4. 其他工具

• K-NN 评估: 运行 scripts/knn 下的脚本。
• 特征可视化 (UMAP): 运行 scripts/umap 下的脚本。
• 目标检测下游任务: 参考 downstream/object_detection/ 目录。

更多详细教程和 API 文档请访问 官方文档。