PRIME: A Few Primitives Can Boost Robustness to Common Corruptions

This is the official repository of PRIME, the data agumentation method introduced in the paper: "PRIME: A Few Primitives Can Boost Robustness to Common Corruptions". PRIME is a generic, plug-n-play data augmentation scheme that consists of simple families of max-entropy image transformations for conferring robustness to common corruptions. PRIME leads to significant improvements in corruption robustness on multiple benchmarks.

Pre-trained models

We provide different models trained with PRIME on CIFAR-10/100 and ImageNet datasets. You can download them from here.

Setup

This code has been tested with Python 3.8.5 and PyTorch 1.9.1. To install required dependencies run:

$ pip install -r requirements.txt

For corruption robustness evaluation, download and extract the CIFAR-10-C, CIFAR-100-C and ImageNet-C datasets from here.

Usage

We provide a script train.py for PRIME training on CIFAR-10/100, ImageNet-100 and ImageNet. For example, to train a ResNet-50 network on ImageNet with PRIME, run:

$ python -u train.py --config=config/imagenet_cfg.py \
    --config.save_dir=<save_dir> \
    --config.data_dir=<data_dir> \
    --config.cc_dir=<common_corr_dir> \
    --config.use_prime=True

Detailed configuration options can be found in config.

Results

Results on ImageNet/ImageNet-100 with a ResNet-50/ResNet-18 (†: without JSD loss)

Results on CIFAR-10/100 with a ResNet-18

Citing this work

@article{PRIME2021,
    title = {PRIME: A Few Primitives Can Boost Robustness to Common Corruptions}, 
    author = {Apostolos Modas and Rahul Rade and Guillermo {Ortiz-Jim\'enez} and Seyed-Mohsen {Moosavi-Dezfooli} and Pascal Frossard},
    year = {2021},
    journal = {arXiv preprint arXiv:2112.13547}
}