Robustness of deep-learning-based compressive sensing

This repository provides code for reproducing the results in the paper:

''Measuring the robustness of deep-learning-based compressive sensing,'' by Mohammad Zalbagi Darestani, Akshay Chaudhari, and Reinhard Heckel

Code by: Mohammad Zalbagi Darestani (mz35@rice.edu) and Reinhard Heckel (rh43@rice.edu)

---

In order to study multiple notions of robustness, the considered problem in the paper is accelerated MRI reconstruction where the task is to reconstruct an image from a few measurements. In this regard, we specifically provide experiments to:

(i) Generate and apply small, adversarial perturbations for un-trained methods like $\mathcal{l}_1$-norm minimization and un-trained neural networks (adversarial_perturbations_untrained.ipynb),

(ii) generate and apply small, adversarial perturbations for trained neural networks like U-net and the end-to-end variational network (VarNet) (adversarial_perturbations_trained.ipynb),

(iii) perform a sample test to apply the models that have only seen samples from the fastMRI dataset to the Stanford set (distribution_shift_fastMRI_to_Stanford.ipynb),

(iv) perform a sample test to apply the models that have only seen knee samples from the fastMRI dataset to a brain image (distribution_shift_fastMRI_knee_to_brain.ipynb),

(v) perform a sample test to apply the models that have only seen samples from the fastMRI dataset to a subset of challenging samples from the fastMRI dataset (distribution_shift_fastMRI_to_adversarially_filtered_data.ipynb),

(vi) put a small feature on the image and recover it using multiple image reconstruction methods (small_features_artificial_fixed_information.ipynb),

(vii) reconstruct an image which contains a real-world pathology and measure how much of that feature is recovered (small_features_real_worl_pathology.ipynb),

(viii) and finally, check whether the choice of basis matters when evaluating the robustness of sparsity-based methods to adversarial perturbations (adversarial_perturbations_basis_doesnt_matter.ipynb).

List of contents

• Setup and installation

• Dataset

• Running the code

• References

• Citation

• License

---

Setup and installation

On a normal computer, it takes aproximately 10 minutes to install all the required softwares and packages.

OS requirements

The code has been tested on the following operating system:

Linux: Ubuntu 16.04.5

Python dependencies

To reproduce the results by running each of the jupyter notebooks, the following softwares are required. Assuming the experiment is being performed in a docker container or a linux machine, the following libraries and packages need to be installed:

    apt-get update
    apt-get install python3.6     # --> or any other system-specific command for installing python3 on your system.
	pip install jupyter
	pip install numpy
	pip install matplotlib
	pip install sigpy
	pip install h5py
	pip install scikit-image
	pip install runstats
	pip install pytorch_msssim
	pip install pytorch-lightning==0.7.5
	pip install test-tube
	pip install Pillow

If pip does not come with the version of python you installed, install pip manually from here. Also, install pytorch from here according to your system specifications.

Install bart toolbox by following the instructions on their home page.

Install PyTorch Wavelets by following these instructions.

Note. After installing pytorch lightning, if you run into a 'state-dict' error for VarNet, you might need to replace parsing.py in /opt/conda/lib/python3.7/site-packages/pytorch_lightning/utilities/parsing.py from here. This is due to the version mismatch in their recent release (0.7.5).

Datasets

The experiments are performed on the following datasets:

The FastMRI dataset.

The Stanford dataset which we collected from mridata.org.

The adversarially-filtered dataset which we created using the invertible Recurrent Inference Machine (i-RIM).

A subset of fastMRI which contains real-world small features. We took the annotations from fastMRI_BB_abnormalities_annotation and created a set of lateral and medical meniscus samples which is available here (see small_features_real_worl_pathology.ipynb to find out how to use this data).

Running the code

You may simply clone this repository and run each notebook to reproduce the results.
Note. You need to download the necessary datasets according to the experiment you intend to run.

References

Code for training the U-net and VarNet is taken from the fastMRI repository.

Code for Deep Decoder is taken from deep_decoder.

Citation

If you use this repository, please cite the original paper:

@inproceedings{darestani2021measuring,
  title={Measuring Robustness in Deep Learning Based Compressive Sensing},
  author={Zalbagi Darestani, Mohammad and S. Chaudhari, Akshay and Heckel, Reinhard},
  booktitle={International Conference on Machine Learning (ICML)},
  year={2021}
}

License

This project is covered by Apache 2.0 License.