Abstract: Optimization of image transformation functions for the purpose of data augmentation has been intensively studied. In particular, adversarial data augmentation strategies, which search augmentation maximizing task loss, show significant improvement in the model generalization for many tasks. However, the existing methods require careful parameter tuning to avoid excessively strong deformations that take away image features critical for acquiring generalization. In this paper, we propose a data augmentation optimization method based on the adversarial strategy called TeachAugment, which can produce informative transformed images to the model without requiring careful tuning by leveraging a teacher model. Specifically, the augmentation is searched so that augmented images are adversarial for the target model and recognizable for the teacher model. We also propose data augmentation using neural networks, which simplifies the search space design and allows for updating of the data augmentation using the gradient method. We show that TeachAugment outperforms existing methods in experiments of image classification, semantic segmentation, and unsupervised representation learning tasks.

TeachAugment: Data Augmentation Optimization Using Teacher Knowledge (CVPR2022, Oral presentation)

Official Implementation of TeachAugment in PyTorch.
arXiv: https://arxiv.org/abs/2202.12513

Requirements

• PyTorch >= 1.9
• Torchvision >= 0.10

Run

Training with single GPU

python main.py --yaml ./config/$DATASET_NAME/$MODEL

Training with single node multi-GPU

python -m torch.distributed.launch --nproc_per_node=$N_GPUS main.py \
    --yaml ./config/$DATASET_NAME/$MODEL --dist

Examples

# Training WRN-28-10 on CIFAR-100
python main.py --yaml ./config/CIFAR100/wrn-28-10.yaml
# Training ResNet-50 on ImageNet with 4 GPUs
python -m torch.distributed.launch --nproc_per_node=4 main.py \
    --yaml ./config/ImageNet/resnet50.yaml --dist

If the computational resources are limited, please try --save_memory option.

Citation

If you find our project useful in your research, please cite it as follows:

@article{suzuki2022TeachAugment
    title={TeachAugment: Data Augmentation Optimization Using Teacher Knowledge},
    author={Suzuki, Teppei},
    journal={arXiv preprint arXiv:2202.12513},
    year={2022},
}

Acknowledgement

The files in ./lib/models and the code in ./lib/augmentation/imagenet_augmentation.py are based on the implementation of Fast AutoAugment.

Download Source Code

Paper Preview