CrIBo

KO | EN

gitlite — search

#python #java #python3 #arduino #golang #machine-learning #rust #html #flask #javascript #seismology #nodejs

★ 12 Open GitHub ↗

This repo contains the code for the ICLR 2024 paper: "CrIBo: Self-Supervised Learning via Cross-Image Object-Level Bootstrapping".

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# CrIBo
[[`OpenReview`](https://openreview.net/pdf?id=3M0GXoUEzP)]

This repo contains the Pytorch implementation of our ICLR 2024 paper (spotlight):
> [**CrIBo: Self-Supervised Learning via Cross-Image Object-Level Bootstrapping**](https://openreview.net/pdf?id=3M0GXoUEzP)
>
>[Tim Lebailly*](https://www.timlebailly.com/), [Thomas Stegmüller*](https://people.epfl.ch/thomas.stegmuller?lang=en), [Behzad Bozorgtabar](https://behzadbozorgtabar.com/), [Jean-Philippe Thiran](https://people.epfl.ch/jean-philippe.thiran) and [Tinne Tuytelaars](https://homes.esat.kuleuven.be/~tuytelaa/).

![alt text](https://raw.githubusercontent.com/tileb1/CrIBo/main/figures/cribo_pipeline.jpg)

## Dependencies
Our code only has a few dependencies. First, install PyTorch for your machine following [https://pytorch.org/get-started/locally/](https://pytorch.org/get-started/locally/).
Then, install other needed dependencies:
```bash
pip install einops
```

## Pretraining
### Single GPU pretraining
Run the [main_cribo.py](main_cribo.py) file. Command line args are defined in [parser.py](cribo_utils/parser.py).
```bash
python main_cribo.py
```

**Make sure to use the right arguments specified in the table below!**

### 1 node pretraining
The code is compatible with slurm. For running on a single node with 8 GPUs:
```bash
#!/bin/bash
#SBATCH --job-name=cribo
#SBATCH --account=<slurm_account>
#SBATCH --cpus-per-task=7
#SBATCH --gpus-per-node=8
#SBATCH --mem=<mem>
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=8
#SBATCH --open-mode=append
#SBATCH --time=4320

master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_ADDR=$master_addr
export MASTER_PORT=12802

srun --unbuffered \
python main_cribo.py \
--imagenet1k_path /path/to/ilsvrc2012 \
--output_dir . \
--n_tokens 32 \
--queue_size 25000 \
--pos_alpha 1.0 1.0 \
--arch vit_small
```
Alternatively, you can use `torchrun` or `torch.distributed.launch` (deprecated):
```bash
python -m torch.distributed.launch --nproc_per_node=8 main_cribo.py \
--imagenet1k_path /path/to/ilsvrc2012 \
--output_dir . \
--n_tokens 32 \
--queue_size 25000 \
--pos_alpha 1.0 1.0 \
--arch vit_small
```
## Citation
If you find our work useful, please consider citing:

```
@inproceedings{
lebailly2024cribo,
title={Cr{IB}o: Self-Supervised Learning via Cross-Image Object-Level Bootstrapping},
author={Tim Lebailly and Thomas Stegm{\"u}ller and Behzad Bozorgtabar and Jean-Philippe Thiran and Tinne Tuytelaars},
booktitle={The Twelfth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=3M0GXoUEzP}
}
```

## Pretrained models
You can download the full checkpoint which contains backbone and projection head weights for both student and teacher networks. We also provide detailed arguments to reproduce our results.

<table class="center">
  <tr>
    <th>pretraining dataset</th>
    <th>arch</th>
    <th>params</th>
    <th>batchsize</th>
    <th>Dense NN retrieval ADE20k (mIoU)</th>
    <th>Dense NN retrieval PVOC12 (mIoU)</th>
    <th colspan="2">download</th>
  </tr>

<tr>
    <th>ImageNet-1k</th>
    <th>ViT-S/16</th>
    <th>21M</th>
    <th>1024</th>
    <th>28.3</th>
    <th>73.2</th>
    <td><a href="https://rdr.kuleuven.be/api/access/datafile/159665">ckpt</a></td>
    <td><a href="https://github.com/tileb1/CrIBo/blob/main/checkpoints/vits16-in_args.json">args</a></td>
  </tr>
  <tr>
    <th>ImageNet-1k</th>
    <th>ViT-B/16</th>
    <th>85M</th>
    <th>1024</th>
    <th>30.0</th>
    <th>74.9</th>
    <td><a href="https://rdr.kuleuven.be/api/access/datafile/159666">ckpt</a></td>
    <td><a href="https://github.com/tileb1/CrIBo/blob/main/checkpoints/vitb16-in_args.json">args</a></td>
  </tr>
</table>

```python
import torch

checkpoint_to_load = 'vitb16-in.pth' # Choose the checkpoint to load
loaded = torch.load(checkpoint_to_load, map_location='cpu')
print(loaded.keys())
```

## Acknowledgments

This code is adapted from [CrOC](https://github.com/stegmuel/CrOC) which is based on the codebase of [DINO](https://github.com/facebookresearch/dino).

CrIBo

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README