Unofficial Pytorch implementation for iMix: A Strategy for regularizing Contrastive Representation Learning.
Paper: https://openreview.net/pdf?id=T6AxtOaWydQ
The environment.yml file contains the required packages. You can install them in a new conda environment as follows:
conda env create -f environment.yml
conda activate iMix
Run an experiment for cifar100 on GPU0:
CUDA_VISIBLE_DEVICES=0 python main.py --save-dir CIFAR100_resnet18 --net resnet18 --dataset cifar100
Resume training
CUDA_VISIBLE_DEVICES=0 python main.py --save-dir CIFAR100_resnet18 --net resnet18 --dataset cifar100 --resume CIFAR100_resnet18/last_model.pth.tar
CUDA_VISIBLE_DEVICES=0 python main.py --save-dir CIFAR100_resnet18 --net resnet18 --dataset cifar100 --no-mix
Multi-gpu is supported using the torch.nn.DataParallel module.
Multiple GPUs on cifar10 using a ResNet50:
CUDA_VISIBLE_DEVICES=0,1 python main.py --save-dir CIFAR10_resnet50 --net resnet50 --dataset cifar10
CUDA_VISIBLE_DEVICES=0 python eval_knn.py CIFAR10_resnet50
CUDA_VISIBLE_DEVICES=0 python eval_linear.py CIFAR10_resnet50
| Dataset | Network | kNN | linear |
|---|---|---|---|
| CIFAR10 | WideResNet28-2 | 81.06 | 81.13 |
| ResNet18 | 91.40 | 93.86 | |
| ResNet50 | 93.40 | 95.25 | |
| CIFAR100 | WideResNet28-2 | 50.77 | 44.30 |
| ResNet18 | 65.28 | 67.69 | |
| ResNet50 | 69.40 | 74.47 |
Define you dataset in the dataset folder using the pytorch templates and the existing files in the datasets folder and add a call in the __init__.py file.
Add you network in the net folder and call it at the beginning of main.py. Don't forget to add a non-linear projection layer for optimal results
@inproceedings{2021_ICLR_iMix,
title="{i-Mix: A Strategy for Regularizing Contrastive Representation Learning}",
author="Lee, Kibok and Zhu, Yian and Sohn, Kihyuk and Li, Chun-Liang and Shin, Jinwoo and Lee, Honglak",
booktitle="{International Conference on Learning Representations (ICLR)}",
year="2021"}