https://github.com/vita-group/augmax
[NeurIPS'21] "AugMax: Adversarial Composition of Random Augmentations for Robust Training" by Haotao Wang, Chaowei Xiao, Jean Kossaifi, Zhiding Yu, Animashree Anandkumar, and Zhangyang Wang.
https://github.com/vita-group/augmax
adversarial-learning imagenet-c model-robustness normalization-techniques out-of-distribution
Last synced: about 1 month ago
JSON representation
[NeurIPS'21] "AugMax: Adversarial Composition of Random Augmentations for Robust Training" by Haotao Wang, Chaowei Xiao, Jean Kossaifi, Zhiding Yu, Animashree Anandkumar, and Zhangyang Wang.
- Host: GitHub
- URL: https://github.com/vita-group/augmax
- Owner: VITA-Group
- License: mit
- Created: 2021-09-28T20:40:58.000Z (over 3 years ago)
- Default Branch: main
- Last Pushed: 2021-12-29T08:00:24.000Z (over 3 years ago)
- Last Synced: 2025-03-29T09:42:03.653Z (about 2 months ago)
- Topics: adversarial-learning, imagenet-c, model-robustness, normalization-techniques, out-of-distribution
- Language: Python
- Homepage:
- Size: 915 KB
- Stars: 125
- Watchers: 8
- Forks: 21
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# AugMax: Adversarial Composition of Random Augmentations for Robust Training
[](https://opensource.org/licenses/MIT)
Haotao Wang, Chaowei Xiao, Jean Kossaifi, Zhiding Yu, Anima Anandkumar, and Zhangyang Wang
In _NeurIPS_ 2021
## Overview
We propose AugMax, a data augmentation framework to unify the diversity and hardness. Being a stronger form of data augmentation, AugMax leads to a significantly augmented input distribution which makes model training more challenging. To solve this problem, we further design a disentangled normalization module, termed DuBIN (Dual-Batch-and-Instance Normalization) that disentangles the instance-wise feature heterogeneity arising from AugMax. AugMax-DuBIN leads to significantly improved out-of-distribution robustness, outperforming prior arts by 3.03%, 3.49%, 1.82% and 0.71% on CIFAR10-C, CIFAR100-C, Tiny ImageNet-C and ImageNet-C.
AugMax achieves a unification between hard and diverse training samples.
AugMax achieves state-fo-the-art performance on CIFAR10-C, CIFAR100-C, Tiny ImageNet-C and ImageNet-C.
## Training
Assume all datasets are stored in ``. For example, CIFAR-10 is in `/cifar-10-batches-py/` and ImageNet training set is in `/imagenet/train`.
AugMax-DuBIN training on `` with `` (The outputs will be saved in ``):
```
python augmax_training_ddp.py --gpu 0 --srp --drp --ds --md --Lambda --steps
```
For example:
AugMax-DuBIN on CIFAR10 with ResNeXt29 (By default we use `Lambda=10` on CIFAR10/100 and Tiny ImageNet.):
```
python augmax_training_ddp.py --gpu 0 --drp /ssd1/haotao/datasets --ds cifar10 --md ResNeXt29 --Lambda 10 --steps 10
```
AugMax-DuBIN on CIFAR100 with ResNet18 (We use `Lambda=1` instead of `Lambda=10` in this particular experiment, as noted in the paper.):
```
python augmax_training_ddp.py --gpu 0 --drp /ssd1/haotao/datasets --ds cifar100 --md ResNet18 --Lambda 1 --steps 10
```
AugMax-DuBIN on ImageNet with ResNet18 (By default we use `Lambda=12` on ImageNet. On ImageNet, **weight decay `wd=1e-4`** instead of the default value in the code, which is for CIFAR-level datasets.):
```
NCCL_P2P_DISABLE=1 python augmax_training_ddp.py --gpu 0 --drp /ssd1/haotao/datasets --ds IN --md ResNet18 --Lambda 12 -e 90 --wd 1e-4 --decay multisteps --de 30 60 --ddp --dist_url tcp://localhost:23456
```
AugMax-DuBIN + DeepAug on ImageNet with ResNet18:
```
NCCL_P2P_DISABLE=1 python augmax_training_ddp.py --gpu 0 --drp /ssd1/haotao/datasets --ds IN --md ResNet18 --deepaug --Lambda 12 -e 30 --wd 1e-4 --decay multisteps --de 10 20 --ddp --dist_url tcp://localhost:23456
```
## Pretrained models
The pretrained models are available on [Google Drive](https://drive.google.com/drive/folders/1GH1fjWQuTYruUU7P7BM52Erg2tAfNJuj?usp=sharing).
## Testing
To test the model trained on `` with `` and saved to ``:
```
python test.py --gpu 0 --ds --drp --md --mode all --ckpt_path
```
For example:
```
python test.py --gpu 0 --ds cifar10 --drp /ssd1/haotao/datasets --md ResNet18_DuBIN --mode all --ckpt_path augmax_training/cifar10/ResNet18_DuBIN/fat-1-untargeted-5-0.1_Lambda10.0_e200-b256_sgd-lr0.1-m0.9-wd0.0005_cos
```
## Citation
```
@inproceedings{wang2021augmax,
title={AugMax: Adversarial Composition of Random Augmentations for Robust Training},
author={Wang, Haotao and Xiao, Chaowei and Kossaifi, Jean and Yu, Zhiding and Anandkumar, Anima and Wang, Zhangyang},
booktitle={NeurIPS},
year={2021}
}
```