Skip to content

Latest commit

 

History

History
76 lines (59 loc) · 4.92 KB

File metadata and controls

76 lines (59 loc) · 4.92 KB

Compressive Visual Representations

This repository contains the source code for our paper, Compressive Visual Representations. We developed information-compressed versions of the SimCLR and BYOL self-supervised learning algorithms, which we call C-SimCLR and C-BYOL, using the Conditional Entropy Bottleneck, and achieved significant improvements in accuracy and robustness, yielding linear evaluation performance competitive with fully supervised models.

cvr_perf

We include implementations of the C-SimCLR and C-BYOL algorithms developed in our paper, as well as SimCLR and BYOL baselines.

Getting Started

Install the necessary dependencies with pip install -r requirements.txt. We recommend creating a new virtual environment.

To train a model with C-SimCLR on ImageNet run bash scripts/csimclr.sh. And to train a model with C-BYOL, run bash scripts/cbyol.sh.

Refer to the scripts for further configuration options, and also to train the corresponding SimCLR and BYOL baselines.

These command lines use the hyperparameters used to train the models in our paper. In particular, we used a batch size of 4096 using 32 Cloud TPUs. Using different accelerators will require changing the batch size. To get started with Google Cloud TPUs, we recommend following this tutorial.

Checkpoints

The following table contains pretrained checkpoints for C-SimCLR, C-BYOL and also their respective baselines, SimCLR and BYOL. All models are trained on ImageNet. The Top-1 accuracy is obtained by training a linear classifier on top of a ``frozen'' backbone, following standard self-supervised learning evaluation protocol.

Algorithm Backbone Training epochs ImageNet Top-1 Checkpoint
SimCLR ResNet 50 1000 70.8 link
SimCLR ResNet 50 2x 1000 74.6 link
C-SimCLR ResNet 50 1000 71.7 link
C-SimCLR ResNet 50 2x 1000 75.0 link
BYOL ResNet 50 1000 74.3 link
BYOL ResNet 50 2x 1000 77.2 link
C-BYOL ResNet 50 1000 75.8 link
C-BYOL ResNet 50 2x 1000 78.9 link
C-BYOL ResNet 101 1000 77.8 link
C-BYOL ResNet 152 1000 78.7 link
C-BYOL ResNet 50 1500 76.0 link

Reference

If you use C-SimCLR or C-BYOL, please use the following BibTeX entry.

@InProceedings{lee2021compressive,
  title={Compressive Visual Representations},
  author={Lee, Kuang-Huei and Arnab, Anurag and Guadarrama, Sergio and Canny, John and Fischer, Ian},
  booktitle={NeurIPS},
  year={2021}
}

Credits

This repository is based on SimCLR. We also match our BYOL implementation in Tensorflow 2 to the original implementation of BYOL in JAX.

Disclaimer: This is not an official Google product.