# Overcoming the Paradox of Certified Training with Gaussian Smoothing

## Evolutional Search
Executes ```evo_train.py```. Calls ```model_wrapper.py``` for computing robust statistics and ```evo_wrapper.py``` for optimizing the model.

## Basic Design

Argument parsing is defined in ```args_factory.py```. It divides the arguments into three groups: **basic** for common options, **train** for training options and **cert** for certifying options. It is recommended to follow this paradigm when adding custom arguments.

The training methods are provided as a special **model wrapper** class. These are all defined in ```model_wrapper.py``` and are subclasses of ```BasicModelWrapper```. These classes are well-defined with inheritance and only overides methods about the corresponding stages. It is recommended to wrap custom methods as a subclass as well.

The main training logic is implemented in ```mix_train.py```. For most cases, trivial modification to this file should be sufficient, e.g., modifying ```save_root``` to adapt to more interested hyperparameter and adding more tracking statistics. It is recommended to follow the comments in the python file rather than place your code arbitrarily.

Tracking statistics of checkpoints is implemented in ```get_stat.py``` in the form of ```xx_loop```, e.g., ```relu_loop``` and ```PI_loop```. These functions are expected to be called at test time and will iterate over the full dataset to compute the corresponding statistics. It is recommended to implement new statistics tracking in similar functional way.

Model certification is done via a combination of IBP (fastest), PGD attack (fast), CROWN-IBP (medium) and MN-BaB (SOTA complete verifier, slow). These are implemented separately in ```mnbab_certify.py``` which should be copied to the root directory of MN-BaB directory (see the top comments in the python file for details). For most cases (except when a new certification method is designed), it is recommended to **not** change this file at all.


## Environments

Recommended environment setup:
```bash
# install miniconda before you start this

# (optional but recommended) install libmambda solver for conda to speed up environment solving
conda update -n base conda
conda install -n base conda-libmamba-solver
conda config --set solver libmamba

# create the environment
conda create --name evo -y python=3.9
conda activate evo

pip install -r requirements.txt
# Dependency evotorch enforces pytorch==2.0.1 automatically, downgrade this manually
pip uninstall torch
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.6 -c pytorch -c nvidia

# install mnbab from public release

```

## Certification

Change the model path in ```scripts/mnbab_certify``` and run it.

## Scripts
In order to train models with PGPE we need init models (we can also start with default init, but it is not recommended since training is very slow).

For training models with back-propagation check ```scripts/train_mnist_grad.sh```. This will save the trained models in ```./SGD_models/```.

For training models with PGPE check ```scripts/train_mnist_pgpe.sh```. This will save the trained models in ```./test_models/```.
For training models with RGS check ```scripts/train_mnist_rgs.sh```. This will save the trained models in ```./test_models/```.

## Models
Released models for MNIST eps=0.1 and CIFAR-10 eps=2/255 CNN5, along with training arguments and certification results can be found in ```./released_models```.
