Awesome

i-RIM for fastMRI

Official implementation of the i-RIM applied to the fastMRI dataset as described in Invert to Learn to Invert and i-RIM applied to the fastMRI challenge. Pre-trained models can be found under Releases.

See some example reconstructions here:

And some numbers:

Installation

To use this code, please run the following commands (preferably in a virtualenv):

git clone --recurse-submodules https://github.com/pputzky/irim_fastMRI.git
cd irim_fastMRI
pip install -r requirements.txt
./install.sh

The above commands will clone this repository with all submodules. Running ./install.sh will install irim as a package in your current Python environment.

Usage

This repository includes two scripts that allow training (scripts.train_model) and running (scripts.run_model) of an i-RIM. Both scripts are derived from the train and run scripts in the fastMRi code base.

To train models as used in our fastMRI challenge submission (see i-RIM applied to the fastMRI challenge), run the following commands (make sure to set $DATA_PATH and $CHECKPOINT_PATH before) :

Single-coil

Train the model

python -m scripts.train_model \
--challenge singlecoil --batch_size 8 --n_steps 8 \
--n_hidden 64 64 64 64 64 64 64 64 64 64 64 64 \
--n_network_hidden 64 64 128 128 256 1024 1024 256 128 128 64 64 \
--dilations 1 1 2 2 4 8 8 4 2 2 1 1 \
--multiplicity 4 --parametric_output \
--loss ssim --resolution 320 --train_resolution 368 368 --lr_gamma 0.1 \
--lr 0.0001 --lr_step_size 30 --num_epochs 50 --optimizer Adam \
--num_workers 8 --report_interval 100 --data_parallel --resume \
--data-path $DATA_PATH --exp_dir $CHECKPOINT_DIR 

Run the model

python -m scripts.run_model --challenge singlecoil --batch-size 8 \
--data-path $DATA_PATH --checkpoint $CHECKPOINT_DIR/best_model.pt \
--out-dir $OUTPUT_DIR --data-split val --mask-kspace

Evaluate reconstructions

python -m external.fastMRI.common.evaluate --challenge singlecoil \
--target-path $DATA_PATH/singlecoil_val/ --predictions-path $OUTPUT_DIR

Multi-coil

python -m scripts.train_model \
--challenge multicoil --batch_size 32 --n_steps 8 \
--n_hidden 96 96 96 96 96 96 96 96 96 96 96 96 \
--n_network_hidden 64 64 128 128 256 1024 1024 256 128 128 64 64 \
--dilations 1 1 2 2 4 8 8 4 2 2 1 1 \
--multiplicity 1 --parametric_output \
--loss ssim --resolution 320 --train_resolution 368 368 --lr_gamma 0.1 \
--lr 0.0001 --lr_step_size 30 --num_epochs 50 --optimizer Adam \
--num_workers 8 --report_interval 100 --data_parallel --resume \
--data-path $DATA_PATH --exp_dir $CHECKPOINT_DIR 

Running and evaluating as above.

RIM baseline on single-coil

python -m scripts.train_model \
--use_rim --challenge singlecoil --batch_size 8 --n_steps 8 \
--loss ssim --resolution 320 --train_resolution 368 368 --lr_gamma 0.1 \
--lr 0.0001 --lr_step_size 30 --num_epochs 50 --optimizer Adam \
--num_workers 8 --report_interval 100 --data_parallel --resume \
--data-path $DATA_PATH --exp_dir $CHECKPOINT_DIR 

Running and evaluating as above.

References

If you use this code or derivatives thereof, please cite the following works

@incollection{pputzky2019,
title = {Invert to Learn to Invert},
author = {Putzky, Patrick and Welling, Max},
booktitle = {Advances in Neural Information Processing Systems 32},
editor = {H. Wallach and H. Larochelle and A. Beygelzimer and F. d\textquotesingle Alch\'{e}-Buc and E. Fox and R. Garnett},
pages = {444--454},
year = {2019},
publisher = {Curran Associates, Inc.},
url = {http://papers.nips.cc/paper/8336-invert-to-learn-to-invert.pdf}
}

@misc{pputzky2019fastMRI,
    title={i-RIM applied to the fastMRI challenge},
    author={Patrick Putzky and Dimitrios Karkalousos and Jonas Teuwen and Nikita Miriakov and Bart Bakker and Matthan Caan and Max Welling},
    year={2019},
    eprint={1910.08952},
    archivePrefix={arXiv},
    primaryClass={eess.IV}
}