This repository contains the code for the paper Modularity Trumps Invariance for Compositonal Robustness.
The dependencies for this project can be found in docker/Dockerfile
. The Docker
image is also currently available on DockerHub.
The code should also be runnable in a conda environment with a fairly standard PyTorch set up using python3.
This is untested, but the key packages should be: pytorch
, torchvision
, numpy
, matplotlib
,
pillow
, pandas
, seaborn
, scikit-image
, opencv-python
and scipy
.
The code in this repo is written to be run on a cluster using Slurm and Singularity. The Docker image is converted into a Singularity image.
The code is almost all executed via sbatch
and the shell scripts in src/scripts
.
- If using Singularity and Slurm you should only need to add your slurm logging path, bind paths and path to the Singularity image in these scripts.
- If working on a system without these tools it should be possible to run the code by taking the commands in these scripts
(beginning
singularity exec ...
) and running only the part of the command that starts withpython ...
.- Commands that contain
$SLURM_ARRAY_TASK_ID
will need to be run with$SLURM_ARRAY_TASK_ID
set to each of the values in the range shown in#SBATCH --array
.
- Commands that contain
All commands should be run from inside the src
directory. The corrupting functions are defined in
data/data_transforms.py
where new corruptions can be added and/or the severity of corruptions can be changed.
The corruptions are designed to work as torchvision transforms
so should be able to be applied to most standard vision datasets.
First download the base dataset to the directory where you want to generate the training data. In an interactive python session run:
from data.emnist import _get_emnist_datasets
_get_emnist_datasets("<path to data root>")
This should create a directory EMNIST/raw
in <path to data root>
containing the base dataset.
To generate the corruptions for training.
Run the first python command in generate_data.sh
after adding the --data-root
flag to match your <path to data root>
.
sbatch scripts/generate_data.sh
CIFAR10 follows a similar process as EMNIST. First download the base data
from data.cifar import _get_cifar_datasets
_get_cifar_datasets("<path to data root>")
You should see the directory cifar-10-batches-py
and file cifar-10-python.tar.gz
in the data root directory.
Then run the second python command in generate_data.sh
. Again changing --data-root
as needed.
sbatch scripts/generate_data.sh
FaceScrub cannot be directly downloaded, you can request access to the dataset here.
We used an internal version of the dataset which is a subset of the original where identities with less than 100 images
are discarded following this paper (section Simulation 2).
We end up with 388 identities (classes). These are in a flat directory structure in <path to data root>/FaceScrub/
with the images named as <identity number>_<image number>.jpg
.
Where <identity number>
is between 100 and 487 (inclusive).
Given this directory structure you can run the third python command in generate_data.sh
. Again changing --data-root
as needed.
sbatch scripts/generate_data.sh
Running the following command will generate a file called corruption_names.pkl
in the data_root/EMNIST
directory.
You will need to change the absolute path to your data root in line 8 before running.
python data/sample_corruptions.py
This file is a random sample to get all the combinations and permuations of corruptions that will be used for testing our models.
CIFAR and FACESCRUB also need a corruption names file. The sampling script can be run again, but easiest is to copy the EMNIST file. In your data root run:
cp EMNIST/corruption_names.pkl CIFAR/
cp EMNIST/corruption_names.pkl FACESCRUB/
After generation you should have in your data root directories called
EMNIST/
CIFAR/
FACESCRUB/
Inside each of these directories you can run the following command to count the number of files
du -a | cut -d/ -f2 | sort | uniq -c | sort -nr
This should output something like this for each dataset (the numbers are important and should match)
To train the different methods we use the scripts scripts/train_monolithic.sh
, scripts/train_modules.sh
and
scripts/train_identity.sh
. The different methods we consider are shown in the above figure, where blue boxes are
trainable parameters and gray boxes are frozen. The monolithic methods (Cross Entropy and Contrastive training) require
one command to run, whereas the modular methods (AutoModules and ImgSpace) require two commands.
To train networks using the cross entropy or contrastive loss run the relevant command in scripts/train_monolithic.sh
.
For example, to train a network on EMNIST using the cross entropy loss run the first line beginning singularity exec ...
which has --experiment "CrossEntropy"
.
Commands are also given for other datasets and for contrastive training. You will need to add the --data-root
flag
to point to your data directory and set --ckpt-path
, --logging-path
and --vis-path
to point to where
you want to save the checkpoints, logs and data visualisations.
To train the modular approach first run the relevant command for the dataset you are using from
scripts/train_identity.sh
. After this network has finished training run the command for the dataset you are using from scripts/train_modules.sh
,
this is the command with --experiment "AutoModules"
. For both commands again set --data-root
,
--ckpt-path
, --logging-path
and --vis-path
.
To train auto-encoders to undo each corruption, first run the relevant command for the dataset you are using from
scripts/train_modules.sh
, this is the command with --experiment "ImgSpace"
. Once the auto-encoders are
trained, two different classifiers can be trained using the command in scripts/train_monolithic.sh
with
--experiment "ImgSpace"
. Again set --data-root
, --ckpt-path
, --logging-path
and --vis-path
.
If the code runs correctly model checkpoints should be saved in the directory specified by --ckpt-path
. The
different methods can then be evaluated by running the scripts/test.sh
script.
Here the arguments --dataset
, --total-n-classes
and --experiment
should be set according to the experiment that you wish to evaluate. The options for --dataset
are EMNIST
,
CIFAR
, and FACESCRUB
which have --total-n-classes
of 47
, 10
and 388
respectively.
The options for --experiment
are CrossEntropy
, Contrastive
, AutoModules
and ImgSpace
.
Once again you will need to point to the directories where the data is stored with --data-root
and where the
checkpoints are stored with --ckpt-path
. The arguments --save-path
, --activations-path
and --vis-path
should be set to the locations where you want to save the results, neural activations and visualisations respectively.
Once the testing completes, box plots comparing experiments can be generated by running python comparison_plots.py
from inside the analysis
directory. Correlations between invariance scores and compositional robustness can be
generated by running python invariance_plots.py
. Flags --data-root
, --results-path
, save-path
and --activations-path
need to be set to the correct directories.
The analysis directory also contains code for saving and processing neural activity for use with the interpretability tool Deephys.
For completeness, we provide the results of our hyper-parameter search here
The training process is as described above but models are evaluated using the code in validate.py
. For two step
methods, we find the best hyper-parameters for the first step, then fix these to find the best hyper-parameters
for the second step.
If you use this code in your research please cite the following paper:
@article{mason2023modularity,
title={Modularity Trumps Invariance for Compositional Robustness},
author={Mason, Ian and Sarkar, Anirban and Sasaki, Tomotake and Boix, Xavier},
journal={arXiv preprint arXiv:2306.09005},
year={2023}
}