Pytorch implementation of SuperMix paper, a supervised method for data augmentation (will appear in CVPR 2021).

• Auguments are:
  • --dataset: specify the dataset, choices: imagenet or cifar100, default: cifar100.
  • --model: specify the supervisor for augmentation. For cifar100, all the models in 'models/__init__.py' can be used. For imagenet, all the models in torchvision.models can be used.
  • --device: specify the device, default: cuda:0.
  • --save_dir: the directory to save the mixed images.
  • --input_dir: the input directory of the imagenet dataset.
  • --bs: batch size, default: 100.
  • --aug_size: number of mixed images to produce, default: 500000.
  • --k: number of input images to be mixed, default: 2.
  • --max_iter: maximum number of iterations on each batch, default: 50.
  • --alpha: alpha value for the Dirichlet distribution, default: 3.
  • --sigma: standard deviation of the Guassian smoothing function, default: 1.
  • --w: spatial size of the mixing masks, default: 8.
  • --lambda_s: multiplier for the sparsity loss, default: 25.
  • --tol: percentage of successfull samples in the batch for early termination, default: 70.
  • --plot: plot the mixed images after generation, default: True

1. Run supermix.py

python3 supermix.py --dataset imagenet --model resnet34 --save_dir ./outputdir --bs 16 --aug_size 50000 --w 16 --sigma 2

2. Sample outputs

1. Download the pretrained model by:

sh scripts/fetch_pretrained_teachers.sh

which saves the models to save/models

2. Run supermix.py

python3 supermix.py --dataset cifar100 --model resnet110 --save_dir ./outputdir --bs 64 --aug_size 50000 --w 8 --sigma 1

3. Sample outputs

Code for the distillation is forked/copied from the official code of CRD

1. Fetch the pretrained teacher models by:

   sh scripts/fetch_pretrained_teachers.sh
   

   which will download and save the models to save/models

2. Produce augmented data using SuperMix by:

   python3 supermix.py --dataset cifar100 --model resnet110 --save_dir ./output --bs 128 --aug_size 500000 --w 8 --sigma 1
   

3. Run the distillation

• using cross-entropy (Equation 9 in the paper) by:

  python3 train_student.py --path_t ./save/models/resnet110_vanilla/ckpt_epoch_240.pth --model_s resnet20 --distill kd --model_s resnet8x4 -r 2.0 -a 0 -b 0 --aug_type supermix --aug_dir ./output --trial 1
  

• using the original distillation objective proposed by Hinton et. al., (Equation 8 in the paper) by:

  python3 train_student.py --path_t ./save/models/resnet110_vanilla/ckpt_epoch_240.pth --model_s resnet20 --distill kd --model_s resnet8x4 -r 1.8 -a 0.2 -b 0 --aug_type supermix --aug_dir ./output --trial 1
  

• where the flags are explained as:

  • --path_t: specify the path of the teacher model
  • --model_s: specify the student model, see 'models/__init__.py' to check the available model types.
  • --distill: specify the distillation method
  • -r: the weight of the cross-entropy loss between logit and ground truth, default: 1
  • -a: the weight of the KD loss, default: None
  • -b: the weight of other distillation losses, default: None
  • --aug_type: type of the augmentation, choices: None, supermix, mixup, cutmix.
  • --aug_dir: the directory of augmented images when supermix is selected for aug_type.
  • --aug_alpha: alpha for the Dirichlet distribution when mixup or cutmix is selected for aug_type.
  • --trial: specify the experimental id to differentiate between multiple runs.

4. (optional) Train teacher networks from scratch. Example commands are in scripts/run_cifar_vanilla.sh

Note: the default setting is for a single-GPU training. If you would like to play this repo with multiple GPUs, you might need to tune the learning rate, which empirically needs to be scaled up linearly with the batch size, see this paper

Performance is measured by classification accuracy (%)

If there is a question regarding any part of the code, or it needs further clarification, please create an issue or send me an email: [email protected].

If you found SuperMix helpful for your research, please cite our paper:

@article{dabouei2020,
  title={SuperMix: Supervising the Mixing Data Augmentation},
  author={Dabouei, Ali and Soleymani, Sobhan and Taherkhani, Fariborz and Nasrabadi, Nasser M},
  journal={arXiv preprint arXiv:2003.05034},
  year={2020}
}

Moreover, if you are developing distillation methods, we encourage you to cite CRD, due to their notable contribution by benchmarking the state-of-the-art methods of distillation.

@inproceedings{tian2019crd,
  title={Contrastive Representation Distillation},
  author={Yonglong Tian and Dilip Krishnan and Phillip Isola},
  booktitle={International Conference on Learning Representations},
  year={2020}
}