This repository contains the code to reproduce the results presented in the experiments.
pip install jax==0.1.75 jaxlib==0.1.52 wfdb
pip install tensorflow tensorflow-probability python_speech_features ujson
pip install tensorflow-datasets seaborn biosppy
We use three datasets: ECG anomaly dataset, Fashion-MNIST and the Speech Command dataset. Fashion-MNIST will be downloaded automatically. You need to download a folder called Resources/ from https://drive.google.com/drive/folders/1EFCC_YKHnQQkgfC9zUsstDldyOtERZ2S?usp=sharing
This folder contains 3 zip files: ecg_recordings.zip needs to be extracted and moved to the ECG folder, speech_dataset.zip needs to be extracted and moved into TensorCommands/. Resources.zip needs to be extracted and moved into the repository.
If you want to run experiments on the cluster, you should move the datasets (Fashion-MNIST,ecg_recordings and Speech Command) to your $SCRATCH directory.
Your $SCRATCH folder should look like this:
ecg_recordings fashion_mnist speech_dataset
But you can change the data directory when you define your experiments.
If you don't want to run the models by defining experiments, you can simply execute the individual scripts. The main scripts are:
main_speech_lsnn.pymain_ecg_lsnn.pymain_CNN_jax.py
You can execute python main_CNN_jax.py --help to view the command line arguments. You can also just run it using python main_CNN_jax.py with the default parameters. The default parameters for each architecture are defined in architectures.py.
Note: This will not save an instance in the database, but it will save the model and the training logs under a specific session id. This way you can use your model. More about how to do this with datajuicer in the following section. You can see the session id in the command line output. When the model is saved, it will say something like: Saved model under XXXXXX_model.json.
We suggest to leave most of the command line arguments untouched and stick to the defaults. The command line arguments that might be of interest include:
- You can specify the dropout probability with
-dropout_prob=0.3 - Change the number of attack steps of our algorith
-n_attack_steps=10 - Change the tradeoff between robustness and performance
-beta_robustness=0.25 - Use AWP
-awp - Use gaussian noise on the foward weights
-noisy_forward_std=0.3 - The relative attack size used during training. This is only relevant if you have
beta_robustnessthat is non-zero.-attack_size_mismatch=0.1 - Change the data directory using this flag
-data_dir=to/your/data
First you need to download the Resources/ folder (see section Data) and place it in this repository. This folder contains the pretrained models, logs and shell logs.
After having setup the environment, you can create the Figures that are in the paper. For the main figures, you can simply run
python make_figures.py
This command will import any experiment that is present in the Experiments folder and execute the visualize() method. This is part of our own experiment manager that we developed for this paper called datajuicer (see more info further down).
Note: You need to have Latex installed for recreating the figures. Follow the steps on https://www.tug.org/texlive/ to install Latex.
We pre-trained the models that were used to run the experiments and re-training all the models takes a lot of time. To retrain all models, you can rename the Sessions folder (rather than deleting it) and simply execute
python train.py
If you want to train the models on a cluster, execute
python train.py -mode=bsub
Note: The command that is needed to execute a script on a cluster varies. To adapt it to your cluster, change the launch_settings dictionary in the architectures.py file.
Note: On the cluster you sometimes need to install the pip packages with the --user option (e.g. pip install --user jax==0.1.75).
If you would like to define a new experiment, create a new experiment file (e.g. my_experiment.py) and move it into the Experiments/ folder. Each experiment is required to implement the train_grid method and a visualize method. The train_grid method defines a model grid that contains the models that you need in order to do the experiment. The following shows an example:
class my_experiment:
@staticmethod
def train_grid():
# Get the default grid (just one element)
grid = [ecg_lsnn.make()]
# For every element in the grid, set the keys in the dict to the specified values
grid = configure(grid, {"boundary_loss":"madry","awp":True})
# For every grid, create 5 copies. Each copy has a different awp_gamma value
grid = split(grid, "awp_gamma", [0.01, 0.1, 0.25, 0.5, 1.0])
return grid
Each element in the grid, which is just a list, is a model that needs to be trained. The grid is defined in such a way that there are no duplicates and only models that are not in the database are retrained.
To train these models, you can now call python train.py -exp=my_experiment -n_threads=5 for local training (add -mode=bsub for cluster training).
The visualizer is used for analysing the results. For example:
@staticmethod
def visualize():
# Get the grid
grid = my_experiment.train_grid()
# Load the models of the grid
grid = run(grid, "train", run_mode="load", store_key="*")("{*}")
grid = configure(grid, {"mode":"direct"})
# Run some function on each model and store the result in an extra key (store_key) in grid
grid = run(grid, get_test_acc, store_key="test_acc")("{*}", "{data_dir}")
After having run some evaluations on the model, one can use query to retrieve results from the grid where certain conditions are met. For example, if we want the test accuracy of the model with awp_gamma=0.1, we run
result = query(grid, "test_acc", where={"awp":True,"awp_gamma":0.1})
For running the experiments we use a custom module called datajuicer which takes care of caching intermediate results and figuring out which data is missing (needs to be calculated) vs which data is already present. Datajuicer manipulates lists of dictionaries which are referred to as grids. The run method takes a grid and a function and returns a decorated function with exactly the same signature except that the inputs are formatted according to the data in the grid. Here is an example where we decorate the inbuilt print function.
import datajuicer as dj
grid = [{"text":"hi", "number":1}, {"text":"hello", "number":2}]
dj.run(grid,print)("{text} {number}")
This outputs:
hi 1
hello 2
and returns (because print returns None):
[None, None]
Instead of specifying a function directly we can also use a key. This is useful if each datapoint in the grid should run the function slightly differently. Here is an example:
import datajuicer as dj
grid = [{"text":"hi", "number":1, "func":(lambda x: print("logging", x))}, {"text":"hello", "number":2, "func":print}]
dj.run(grid,"func")("{text} {number}")
This outputs:
logging hi 1
hello 2
As a special 'key', * gets replaced by the entire dictionary.
By default, the decorated function returns a list of all the values returned by the original function, but we can change this by passing `store_key' to the run function. If store_key is set to * then it is assumed that the function returns a dictionary and this dictionary is merged with each data in the grid. Otherwise, if the store_key is set to some string, whatever the function returns is stored in the store_key and the decorated function returns a copy of the original grid that also contains the returned values.
The formatting can get arbitrarily deep and even call functions. It is accessable via the format_template function. Here is an example:
import datajuicer as dj
grid = [{"mood":"good", "topic":"food","fav_food":"cheese"}, {"mood":"lousy", "topic":"drink", "fav_drink":"juice"}]
def say(mood, msg):
if mood == "good":
return "Today is a great day! I just want to say: " + msg
else:
return msg
grid = dj.configure(grid, {"say": say, "favorites":"{say({mood},My favorite {topic} is {fav_{topic}}.)}"})
dj.run(grid, print)("{favorites}")
This outputs:
Today is a great day! I just want to say: My favorite food is cheese.
My favorite drink is juice.
One important feature of datajuicer is the @cachable() decorator. Whenever you declare a function, you can add @cachable() to the line above and then whenever this function is called using run it will not recalculate for inputs it has already seen. Here is an example:
import datajuicer as dj
@dj.cachable()
def add(a, b):
print("calculating...")
return a + b
grid = [{}]
grid = dj.split(grid, "a", [1,2,3])
grid = dj.split(grid, "b", [10,20,30])
print(dj.run(grid, add)("{a}", "{b}"))
grid.append({"a":1, "b":1})
print(dj.run(grid, add)("{a}", "{b}"))
You can pass a custom saver or loader to cachable if you do not want to use the defaults. Also if some of the inputs to the function will not have an effect on the output, you can ignore them by passing a list of strings to cachable with the keyword dependencies. You can also use the syntax arg_name:key_name if you pass a dictionary to a function and the output depends on a specific key of that dictionary. For example:
@cachable(dependencies=["a", "b", "my_dict:my_key"])
def foo(a, b, my_dict, log_directory):
pass
In this example, the runner will not recalculate just because a different log_directory is used.