Awesome

Mixed Barlow Twins for Self-Supervised Representation Learning

Guarding Barlow Twins Against Overfitting with Mixed Samples<br>

Wele Gedara Chaminda Bandara (Johns Hopkins University), Celso M. De Melo (U.S. Army Research Laboratory), and Vishal M. Patel (Johns Hopkins University) <br>

1 Overview of Mixed Barlow Twins

TL;DR

• Mixed Barlow Twins aims to improve sample interaction during Barlow Twins training via linearly interpolated samples.
• We introduce an additional regularization term to the original Barlow Twins objective, assuming linear interpolation in the input space translates to linearly interpolated features in the feature space.
• Pre-training with this regularization effectively mitigates feature overfitting and further enhances the downstream performance on CIFAR-10, CIFAR-100, TinyImageNet, STL-10, and ImageNet datasets.

$C^{MA} = (Z^M)^TZ^A$

$C^{MB} = (Z^M)^TZ^B$

$C^{MA}_{gt} = \lambda (Z^A)^TZ^A + (1-\lambda)\mathtt{Shuffle}^*(Z^B)^TZ^A$

$C^{MB}_{gt} = \lambda (Z^A)^TZ^B + (1-\lambda)\mathtt{Shuffle}^*(Z^B)^TZ^B$

2 Usage

2.1 Requirements

Before using this repository, make sure you have the following prerequisites installed:

• Anaconda
• PyTorch

You can install PyTorch with the following command (in Linux OS):

conda install pytorch torchvision pytorch-cuda=11.8 -c pytorch -c nvidia

2.2 Installation

To get started, clone this repository:

git clone https://github.com/wgcban/mix-bt.git

Next, create the conda environment named ssl-aug by executing the following command:

conda env create -f environment.yml

All the train-val-test statistics will be automatically upload to wandb, and please refer wandb-quick-start documentation if you are not familiar with using wandb.

2.3 Supported Pre-training Datasets

This repository supports the following pre-training datasets:

• CIFAR-10: https://www.cs.toronto.edu/~kriz/cifar.html
• CIFAR-100: https://www.cs.toronto.edu/~kriz/cifar.html
• Tiny-ImageNet: https://github.com/rmccorm4/Tiny-Imagenet-200
• STL-10: https://cs.stanford.edu/~acoates/stl10/
• ImageNet: https://www.image-net.org

CIFAR-10, CIFAR-100, and STL-10 datasets are directly available in PyTorch.

To use TinyImageNet, please follow the preprocessing instructions provided in the TinyImageNet-Script. Download these datasets and place them in the data directory.

2.4 Supported Transfer Learning Datasets

You can download and place transfer learning datasets under their respective paths, such as 'data/DTD'. The supported transfer learning datasets include:

• DTD: https://www.robots.ox.ac.uk/~vgg/data/dtd/
• MNIST: http://yann.lecun.com/exdb/mnist/
• FashionMNIST: https://github.com/zalandoresearch/fashion-mnist
• CUBirds: http://www.vision.caltech.edu/visipedia/CUB-200-2011.html
• VGGFlower: https://www.robots.ox.ac.uk/~vgg/data/flowers/102/
• Traffic Signs: https://benchmark.ini.rub.de/gtsdb_dataset.html
• Aircraft: https://www.robots.ox.ac.uk/~vgg/data/fgvc-aircraft/

2.5 Supported SSL Methods

This repository supports the following Self-Supervised Learning (SSL) methods:

• SimCLR: contrastive learning for SSL
• BYOL: distilation for SSL
• Witening MSE: infomax for SSL
• Barlow Twins: infomax for SSL
• Mixed Barlow Twins (ours): infomax + mixed samples for SSL

2.6 Pre-Training with Mixed Barlow Twins

To start pre-training and obtain k-NN evaluation results for Mixed Barlow Twins on CIFAR-10, CIFAR-100, TinyImageNet, and STL-10 with ResNet-18/50 backbones, please run:

sh scripts-pretrain-resnet18/[dataset].sh

sh scripts-pretrain-resnet50/[dataset].sh

To start the pre-training on ImageNet with ResNet-50 backbone, please run:

sh scripts-pretrain-resnet18/imagenet.sh

2.7 Linear Evaluation of Pre-trained Models

Before running linear evaluation, ensure that you specify the model_path argument correctly in the corresponding .sh file.

To obtain linear evaluation results on CIFAR-10, CIFAR-100, TinyImageNet, STL-10 with ResNet-18/50 backbones, please run:

sh scripts-linear-resnet18/[dataset].sh

sh scripts-linear-resnet50/[dataset].sh

To obtain linear evaluation results on ImageNet with ResNet-50 backbone, please run:

sh scripts-linear-resnet50/imagenet_sup.sh

2.8 Transfer Learning of Pre-trained Models

To perform transfer learning from pre-trained models on CIFAR-10, CIFAR-100, and STL-10 to fine-grained classification datasets, execute the following command, making sure to specify the model_path argument correctly:

sh scripts-transfer-resnet18/[dataset]-to-x.sh

3 Pre-Trained Checkpoints

Download the pre-trained models from GitHub (Releases v1.0.0) and store them in checkpoints/. This repository provides pre-trained checkpoints for both ResNet-18 and ResNet-50 architectures.

3.1 ResNet-18 [CIFAR-10, CIFAR-100, TinyImageNet, and STL-10]

3.2 ResNet-50 [CIFAR-10, CIFAR-100, TinyImageNet, and STL-10]

3.3. ResNet-50 on ImageNet (300 epochs)

Setting: epochs = 300, $d$ = 8192, $\lambda_{BT}$ = 0.0051

3.4. ResNet-50 on ImageNet (1000 epochs)

Setting: epochs = 1000, $d$ = 8192, $\lambda_{BT}$ = 0.0051, $\lambda_{reg}$=2.0

4 Training/Val Logs

3.1 Pre-trianing for 300 epochs

Logs are available on wandb and can access via following links:

• imagenet pre-training: https://api.wandb.ai/links/cha-yas/5olb2sar
• imagenet linear probing: https://api.wandb.ai/links/cha-yas/9tb0ksfp

Here we provide some training and validation (linear probing) statistics for Barlow Twins vs. Mixed Barlow Twins with ResNet-50 backbone on ImageNet:

<img src="figs/in-loss-bt.png" width="256"/> <img src="figs/in-loss-reg.png" width="256"/> <img src="figs/in-linear.png" width="256"/>

3.1 Pre-trianing for 1000 epochs

We also provide trianing-val statistics for our pre-trained model for 1000 epochs. <img src="figs/in-loss-bt-1000e.png" width="256"/> <img src="figs/in-loss-reg-1000e.png" width="256"/> <img src="figs/in-linear-1000e.png" width="256"/>

:fire: Access pre-training statistcis on wandb: wandb-imagenet-pretrain

5 Disclaimer

A large portion of the code is from Barlow Twins HSIC (for experiments on small datasets: CIFAR-10, CIFAR-100, TinyImageNet, and STL-10) and official implementation of Barlow Twins here (for experiments on ImageNet), which is a great resource for academic development.

Also, note that the implementation of SOTA methods (SimCLR, BYOL, and Witening-MSE) in ssl-sota are copied from Witening-MSE.

We would like to thank all of them for making their repositories publicly available for the research community. 🙏

6 Reference

If you feel our work is useful, please consider citing our work. Thanks!

@misc{bandara2023guarding,
      title={Guarding Barlow Twins Against Overfitting with Mixed Samples}, 
      author={Wele Gedara Chaminda Bandara and Celso M. De Melo and Vishal M. Patel},
      year={2023},
      eprint={2312.02151},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

7 License

This code is under MIT licence, you can find the complete file here.