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
对backbone进行微调,mixup、cutmix和random augment。 - [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
- Supervised Contrastive Learning
- SwAV
- VIbCReg
- VICReg
- W-MSE
额外特性
骨干网络
数据
- 使用Nvidia Dali将数据处理速度提高了多达100%。
- 灵活的增强。
评估
- 通过停止梯度的在线线性评估,用于更容易的调试和原型设计(动量骨干网络也可选用)。
- 标准离线线性评估。
- 在线和离线K-NN评估。
- 使用UMAP自动特征空间可视化。
训练技巧
- PyTorch Lightning的所有优点(混合精度、梯度累积、裁剪等等)。
- 通道最后转换
- 遵循SwAV的多裁剪数据加载:
- 注意:目前只有SimCLR、BYOL和SwAV支持这一功能。
- 从权重衰减和LARS中排除批量归一化和偏置。
- 投影头没有LR调度器(如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 install
**注意:**如果您在安装dali时遇到问题,请按照他们的指南进行安装。
**注意2:**考虑安装Pillow-SIMD以获得更好的加载时间(当不使用Dali时)。
**注意3:**即将在pip上发布。
训练
要预训练骨干网络,请参照scripts/pretrain/
中的多个bash文件。
我们现在使用Hydra来处理配置文件,所以常见的语法如下:
python3 main_pretrain.py \
# 训练脚本文件夹的路径
--config-path scripts/pretrain/imagenet-100/ \
# 训练配置名称
--config-name barlow.yaml
# 添加新参数(例如那些未在yaml文件中定义的)
# 通过执行++new_argument=VALUE
# pytorch lightning的参数也可以在这里添加。
之后,对于离线线性评估,请参照scripts/linear
中的示例,或者对于微调整个骨干网络,请参照scripts/finetune
。
对于k-NN评估和UMAP可视化,请查看scripts/{knn,umap}
中的脚本。
**注意:**文件尽量保持最新,并尽可能遵循每篇论文推荐的参数,但在运行之前请检查它们。
教程
请查看我们的文档和教程:
- 概述
- 离线线性评估
- [对象检测](https://github.com/vt | 方法 | 骨干网络 | 轮数 | 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
方法 | 骨干网络 | 轮数 | 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 | Acc@1 (在线) | Acc@1 (离线) | Acc@5 (在线) | Acc@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 | 待定 | 87.72 | 待定 | :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 | 待定 | 94.06 | 待定 | :link: |
Simsiam | ResNet18 | 400 | :heavy_check_mark: | 74.54 | 78.72 | 93.16 | 94.78 | :link: |
SupCon | ResNet18 | 400 | :heavy_check_mark: | 84.40 | 待定 | 95.72 | 待定 | :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 | Acc@1 (在线) | Acc@1 (离线) | Acc@5 (在线) | Acc@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的训练效率
我们报告了一些方法在配备英特尔i9-9820X和两个RTX2080ti的服务器上使用ResNet18,有无DALI(每个GPU 4个工作进程)的训练效率。
方法 | Dali | 20轮的总时间 | 1轮的时间 | 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 and Enrico Fini and Moin Nabi and Nicu Sebe and Elisa Ricci},
title = {solo-learn: A Library of Self-supervised Methods for Visual Representation Learning},
journal = {Journal of Machine Learning Research},
year = {2022},
volume = {23},
number = {56},
pages = {1-6},
url = {http://jmlr.org/papers/v23/21-1155.html}
}