Mind Your Augmentation: The Key to De-coupling Dense Self-supervised Learning

Code and pre-trained models of ICLR 2024 paper Mind Your Augmentation: The Key to De-coupling Dense Self-supervised Learning.

[paper] [project page]

The work presents a thorough investigation of the impact of coupling shortcuts in dense self-supervised learning. We propose a de-coupling strategy to integrate with the existing dense self-supervised learning methods to mitigate the negative impact of the coupling shortcuts, which consistently improves the performance of dense-level evaluation tasks.

Acknowledgement: Our project is built using the iBOT repository.

News

June 2024—Release the iBOT-based pre-trained models and code.

Todo

Release the code and pre-trained models for other pre-training methods.

Installation

See installation structions for details. You also need to install Kornia for data augmentation:

pip install kornia==0.6.9

Data Preparation

We perform pre-training on MS-COCO and dense-level evaluation on MS-COCO and ADE20K. Please download the datasets from their official websites and organize the data as follows:

 ── your project path/
    └── data/
        └── ade
        │   ├── ADEChallengeData2016/
        │   │   ├── annotations/
        │   │   ├── images/
        └── coco
            ├── annotations/
            ├── train2017/
            └── val2017/

Pre-training

iBOT Pre-Training with ViTs

We use an 8-GPU machine with RTX 3090 GPUs to pre-train the models. To pre-train a ViT-S/16 model on MS-COCO with DDP, you can directly run the script as follows:

bash configs/Pretrain.sh

The parameter --dc is used to enable the de-coupling strategy. You can remove it to pre-train the model with the vanilla setting.

Pre-Trained Models

The pre-trained models are available for download.DC denotes the de-coupling strategy.

Model	Arch.	DC	COCO Det.	COCO ISeg.	ADE20K Seg.	Checkpoint
iBOT	ViT-S/16	×	42.3	37.0	39.9	ckpt
iBOT	ViT-S/16	√	45.1	39.1	41.6	ckpt

Downstream Evaluation

The evaluation protocol is consistent with iBOT. You can refer to Evaluating iBOT on Downstream Tasks for details. Or you can directly run the script to perform dense-level evaluation on MS-COCO and ADE20K.

License

This repository is released under the MIT license as found in the LICENSE file.

Citation

If you find this repository useful, please consider giving a star ⭐ and citation:

@inproceedings{
qiu2024mind,
title={Mind Your Augmentation: The Key to Decoupling Dense Self-Supervised Learning},
author={Congpei Qiu and Tong Zhang and Yanhao Wu and Wei Ke and Mathieu Salzmann and Sabine S{\"u}sstrunk},
booktitle={The Twelfth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=WQYHbr36Fo}
}

ztt1024/denseSSL