Awesome

"Make your AI Robust."

MAIR is a PyTorch-based adversarial training framework. MAIR aims to (1) provide easy implementation of adversarial training methods and (2) make it easier to evaluate the adversarial robustness of deep learning models.

Adversarial training has become the de facto standard method for improving the robustness of models against adversarial examples. However, while writing our paper "Fantastic Robustness Measures: The Secrets of Robust Generalization" [Paper, Article], we realized that there is no framework integrating adversarial training methods. Therefore, to promote reproducibility and transparency in deep learning, we integrated the algorithms, tools, and pre-trained models.

Citation: If you use this package, please cite the following BibTex (GoogleScholar):

@inproceedings{
    kim2023fantastic,
    title={Fantastic Robustness Measures: The Secrets of Robust Generalization},
    author={Hoki Kim and Jinseong Park and Yujin Choi and Jaewook Lee},
    booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
    year={2023},
    url={https://openreview.net/forum?id=AGVBqJuL0T}
}

Benchmarks on adversarially trained models are available at our notion.

Installation and usage

Installation

pip install git+https://github.com/Harry24k/MAIR.git

Usage

import mair

How to train a model?

Step1. Load model as follows:

model = ...
rmodel = mair.RobModel(model, n_classes=10).cuda()

Step2. Set trainer as follows:

from mair.defenses import AT
# Set adversarial training method: [Strandard, AT, TRADES, MART].
trainer = AT(rmodel, eps=EPS, alpha=ALPHA, steps=STEPS)
# Set recording information.
trainer.record_rob(train_loader, val_loader, eps=EPS, alpha=2/255, steps=10, std=0.1)
# Set detail training methods.
trainer.setup(optimizer="SGD(lr=0.1, momentum=0.9)",
              scheduler="Step(milestones=[100, 150], gamma=0.1)",
              scheduler_type="Epoch",
              minimizer=None, # or "AWP(rho=5e-3)",
              n_epochs=200
             )

Step3. Fit model as follows:

trainer.fit(train_loader=train_loader,
            n_epochs=200,
            save_path='./models/', 
            save_best={"Clean(Val)":"HBO", "PGD(Val)":"HB"},
            # 'save_best': model with high PGD are chosen, 
            # while in similar cases, model with high Clean are selected.
            save_type="Epoch", 
            save_overwrite=False, 
            record_type="Epoch"
           )

How to evaluate a model?

Step1. Transform model as follows:

model = ...
rmodel = mair.RobModel(model, n_classes=10).cuda()

Step2. Evaluate model as follows:

rmodel.eval_accuracy(test_loader)  # clean accuracy
rmodel.eval_rob_accuracy_gn(test_loader)  # gaussian noise accuracy
rmodel.eval_rob_accuracy_fgsm(test_loader, eps)  # FGSM accuracy
rmodel.eval_rob_accuracy_pgd(test_loader, eps, alpha, steps)  # PGD accuracy

Please refer to demo for details.

Adversarial Benchmarks & Pre-trained models

Here is our (selected) benchmark on popular techniques in adversarial training frameworks.

Note that all robustness (or robust accuracy) are measured on CIFAR-10 test dataset. Through our notion, you can check more detailed benchmarks.

ResNet

WRN28-10

WRN34-10

Based on our notion, we built a hub module to support the direct use of our pretrained models.

from mair.hub import load_pretrained
rmodel = load_pretrained("CIFAR10_ResNet18_AT(eps=8, alpha=2, steps=10)", flag='Best', save_dir="./")

Please refer to demo for details.

Or you can use Google-drive.

In each folder, we upload four different files:

• log.txt: training log during training.
• last.pth: model at the end of epoch.
• init.pth: model at the start of epoch.
• best.pth: best model selected by the argment save_best in trainer.fit.

To load model,

rmodel.load_dict('./models/.../best.pth')

We are excited to share modes with the community, but we've run into a storage limitation on Google Drive. Any help would be greatly appreciated!

Contribution

We welcome all contribution to MAIR in many forms :smiley:. Especially, we are looking for diverse adversarial training methods beyond AT, TRADES, and MART.

Future work

• Merge measures.
• Generalize attacks gathered from torchattacks.
• ...