Robust Training for Image Classification
========================================

Overview
--------

This project contains the code for adversarial training on classification models. Heavy augmentations can be used in adversarial training. The well-known `timm <https://github.com/huggingface/pytorch-image-models>`_ project is used as the default classification library. 

Preparation 
-----------

- Dataset  

  We use ``ImageNet Validation Set`` as the default dataset to evaluate adversarial robustness of classification models. Please download `ImageNet <https://www.image-net.org/>`_ dataset first. If you want to use your own datasets, please define their ``torch.utils.data.Dataset`` class and corresponding ``transform``.

- Classification Models 
 
  To build a image classification model, you can create a model class from `timm <https://github.com/huggingface/pytorch-image-models>`_ library or you can define custom network of ``torch.nn.Module``.

Robust Training
---------------

Before start, modify the default attacking config files in ``train_configs``. It contains the augmentation parameters, dataset directory, etc.

Then, you can run the following command to start training ResNet50.

.. code-block:: bash

   cd robust_training
   python -m torch.distributed.launch --nproc_per_node=<num-of-gpus-to-use> adversarial_training.py --configs=./train_configs/resnet50.yaml