This repository contains the code used in the experiments of Rommel, Moreau and Gramfort, "Deep invariant networks with differentiable augmentation layers", NeurIPS, 2022.
Start by cloning the repository:
git clone https://github.com/cedricrommel/augnet.git
This work depends on the following packages:
- pytorch>=1.8.0
- torchvision>=0.11.1
- kornia>=0.6.3
- scikit-learn
- numpy
- pandas
- scipy
- matplotlib
- seaborn
- pytest
The EEG experiments (not yet available) further require:
- MNE Python
- Skorch
- Braindecode
All depedencies can be installed into a new environment using Anaconda by running:
conda env create -f environment.yml
conda activate augnet
from the package root directory.
After installing the necessary requirements, you should install the package itself. For this, go to the root directory of the package and run
pip install -e .
You can create augmentation layers using the augnet.models.AugmentationLayer
class, by passing to it a list of differentiable augmentations, such as those
in the module augnet.augmentations.vision
:
from augnet.augmentations.vision import (
TranslateX,
TranslateY,
ShearX,
ShearY,
Rotate,
)
from augnet.models import AugmentationLayer
n_layers = 3
init_mag = 0.1
fixed_prob = 1.0
aug_layers = list()
for _ in range(n_layers):
# create augmentations
augmentations = [
TranslateX(
initial_magnitude=init_mag,
initial_probability=fixed_prob,
),
TranslateY(
initial_magnitude=init_mag,
initial_probability=fixed_prob
),
Rotate(
initial_magnitude=init_mag,
initial_probability=fixed_prob,
magnitude_scale=180
),
ShearX(
initial_magnitude=init_mag,
initial_probability=fixed_prob
),
ShearY(
initial_magnitude=init_mag,
initial_probability=fixed_prob
),
]
# create layers
aug_layer = AugmentationLayer(
augmentations=augmentations,
temperature=0.01,
)
aug_layers.append(aug_layer)
A list of augmentation layers can be grouped into an AugmentationModule
as follows:
from augnet.models import AugmentationModule
aug_module = AugmentationModule(aug_layers)
The latter can be used to wrap any trunk model using the AugAveragedModel
class, in order to create an AugNet meta-model:
from augnet.models import AugAveragedModel, SimpleMLP
trunk_model = SimpleMLP(
n_neurons=10,
n_layers=3,
in_shape=32*32*3,
num_classes=4,
)
# create augnet
augnet = AugAveragedModel(
model=trunk_model,
aug=aug_module,
ncopies=4,
)
The obtained model can be trained seemlessly as any other pytorch model.
Mario and Iggy experiment code (section 4.1) can be found in
./experiments/mario_iggy
. To reproduce the paper experiments, just move to
that folder and execute
bash run_experiments.sh
Figures will be plotted and saved into ./experiments/mario_iggy/figures/
,
while outputs used to create them are stored in
./experiments/mario_iggy/saved-outputs
.
Sinusoids experiments (section 4.2 and 4.3) can be found in
./experiments/sinusoids
. To reproduce the paper experiments, just move to
that folder and execute
bash run_experiments.sh
Figures will be plotted and saved into ./experiments/sinusoids/figures/
,
while outputs used to create them are stored in
./experiments/sinusoids/saved-outputs
.
CIFAR10 experiments code (section 5.1) can be found in ./experiments/cifar10
.
To reproduce the baselines from the paper, move there and execute
bash run_experiments.sh $DEVICE
where $DEVICE
is the device you want to run the training on (e.g. cuda:0
).
Training can take a lot of time though as they are not parallelized here (~45h for the baselines and ~20h for Augnet).
Figures will be plotted and saved into ./experiments/cifar10/figures/
, while outputs will be stored in ./experiments/cifar10/saved-outputs
.
MASS experiments code (section 5.2) can be found in ./experiments/mass
.
Unfortunately, this experiment still cannot be reproduced since it relies on a package that was not yet publicly released. This repository will be updated as soon as the situation changes.
If you use the code in this repository, please consider citing:
@inproceedings{rommel2022deep,
title={Deep invariant networks with differentiable augmentation layers},
author={Rommel, C\'edric and Thomas Moreau and Alexandre Gramfort},
year={2022},
booktitle={{Advances in Neural Processing Systems (NeurIPS)}},
}
The differentiable computer vision augmentations in this repository were copied and modified from https://github.com/moskomule/dda and the Augerino model together with a few other pieces are strongly inspired by https://github.com/g-benton/learning-invariances.