#### NOTE
New to DeepInverse? Get started with the basics with the
[5 minute quickstart tutorial](https://deepinv.org/auto_examples/basics/demo_quickstart.html.md#sphx-glr-auto-examples-basics-demo-quickstart-py)..
# Self-supervised learning with Equivariant Imaging for MRI.
This example shows you how to train a reconstruction network for an MRI inverse problem on a fully self-supervised way, i.e., using measurement data only.
The equivariant imaging loss is presented in Chen *et al.*[1](#footcite-chen2021equivariant).
```Python
from pathlib import Path
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
import deepinv as dinv
from deepinv.datasets import SimpleFastMRISliceDataset
from deepinv.utils import load_degradation
from deepinv.models.utils import get_weights_url
from deepinv.models import MoDL
```
## Setup paths for data loading and results.
```Python
BASE_DIR = Path(".")
DATA_DIR = BASE_DIR / "measurements"
CKPT_DIR = BASE_DIR / "ckpts"
# Set the global random seed from pytorch to ensure reproducibility of the example.
torch.manual_seed(0)
device = dinv.utils.get_device()
```
```none
Selected GPU 0 with 22453.25 MiB free memory
```
## Load base image datasets and degradation operators.
In this example, we use a mini demo subset of the single-coil [FastMRI dataset](https://fastmri.org/)
as the base image dataset, consisting of 2 knee images of size 320x320.
#### SEE ALSO
Datasets [`deepinv.datasets.FastMRISliceDataset`](https://deepinv.org/api/stubs/deepinv.datasets.FastMRISliceDataset.html.md#deepinv.datasets.FastMRISliceDataset) [`deepinv.datasets.SimpleFastMRISliceDataset`](https://deepinv.org/api/stubs/deepinv.datasets.SimpleFastMRISliceDataset.html.md#deepinv.datasets.SimpleFastMRISliceDataset)
: We provide convenient datasets to easily load both raw and reconstructed FastMRI images.
You can download more data on the [FastMRI site](https://fastmri.med.nyu.edu/).
#### IMPORTANT
By using this dataset, you confirm that you have agreed to and signed the [FastMRI data use agreement](https://fastmri.med.nyu.edu/).
#### NOTE
We reduce to the size to 128x128 for faster training in the demo.
```Python
operation = "MRI"
img_size = 128
transform = transforms.Compose([transforms.Resize(img_size)])
train_dataset = SimpleFastMRISliceDataset(
transform=transform, train_percent=0.5, train=True, download=True
)
test_dataset = SimpleFastMRISliceDataset(
transform=transform, train_percent=0.5, train=False
)
```
```none
/local/jtachell/deepinv/deepinv/deepinv/datasets/fastmri.py:109: DeprecationWarning: The tuple format for dataset outputs is deprecated and will be removed in a future version.It is recommended to set `use_dict_output=True` for better readability and flexibility in returned outputs.The default is currently `False` for backward compatibility, but will be switched to `True` in a future version.
super().__init__(use_dict_output=use_dict_output)
```
## Generate a dataset of knee images and load it.
```Python
# training the model from scratch can take a lot of time, so here we load pre-trained weights
# if you want to train the model from scratch, set this to True
train_from_scratch = False
mask = load_degradation("mri_mask_128x128.npy")
# defined physics
physics = dinv.physics.MRI(mask=mask, device=device)
# Use parallel dataloader if using a GPU to speed up training,
# otherwise, as all computes are on CPU, use synchronous data loading.
num_workers = 4 if torch.cuda.is_available() else 0
# number of images used for training
n_images_max = 900 if train_from_scratch else 16
my_dataset_name = "demo_equivariant_imaging"
measurement_dir = DATA_DIR / "fastmri" / operation
deepinv_datasets_path = dinv.datasets.generate_dataset(
train_dataset=train_dataset,
test_dataset=test_dataset,
physics=physics,
device=device,
save_dir=measurement_dir,
train_datapoints=n_images_max,
num_workers=num_workers,
dataset_filename=str(my_dataset_name),
)
train_dataset = dinv.datasets.HDF5Dataset(
path=deepinv_datasets_path, train=True, use_dict_output=True
)
test_dataset = dinv.datasets.HDF5Dataset(
path=deepinv_datasets_path, train=False, use_dict_output=True
)
```
```none
/local/jtachell/deepinv/deepinv/.pixi/envs/docs/lib/python3.12/multiprocessing/popen_fork.py:66: DeprecationWarning: This process (pid=4088353) is multi-threaded, use of fork() may lead to deadlocks in the child.
self.pid = os.fork()
Dataset has been saved at measurements/fastmri/MRI/demo_equivariant_imaging0.h5
```
## Set up the reconstruction network
As a (static) reconstruction network, we use an unrolled network
(half-quadratic splitting) with a trainable denoising prior based on the
DnCNN architecture which was proposed in MoDL Aggarwal *et al.*[2](#footcite-aggarwal2018modl).
See [`deepinv.models.MoDL`](https://deepinv.org/api/stubs/deepinv.models.MoDL.html.md#deepinv.models.MoDL) for details.
```Python
model = MoDL().to(device)
```
## Set up the training parameters
We choose a self-supervised training scheme with two losses: the measurement consistency loss (MC)
and the equivariant imaging loss (EI).
The EI loss requires a group of transformations to be defined. The forward model should not be equivariant to
these transformations Tachella *et al.*[3](#footcite-tachella2023sensing).
Here we use the group of 4 rotations of 90 degrees, as the accelerated MRI acquisition is
not equivariant to rotations (while it is equivariant to translations).
See [docs](https://deepinv.org/user_guide/training/transforms.html.md#transform) for full list of available transforms.
#### NOTE
We use a pretrained model to reduce training time. You can get the same results by training from scratch
for 150 epochs using a larger knee dataset of ~1000 images.
```Python
epochs = 150 if train_from_scratch else 1
learning_rate = 5e-4
batch_size = 16 if torch.cuda.is_available() else 1
# choose self-supervised training losses
# generates 4 random rotations per image in the batch
losses = [dinv.loss.MCLoss(), dinv.loss.EILoss(dinv.transform.Rotate(n_trans=4))]
# choose optimizer and scheduler
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-8)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(epochs * 0.8) + 1)
if not train_from_scratch:
# start with a pretrained model to reduce training time
file_name = "new_demo_ei_ckp_150_v3.pth"
url = get_weights_url(model_name="demo", file_name=file_name)
ckpt = torch.hub.load_state_dict_from_url(
url,
map_location=lambda storage, loc: storage,
file_name=file_name,
)
# load a checkpoint to reduce training time
model.load_state_dict(ckpt["state_dict"])
optimizer.load_state_dict(ckpt["optimizer"])
```
```none
/local/jtachell/deepinv/deepinv/deepinv/transform/rotate.py:49: UserWarning: The default interpolation mode will be changed to bilinear interpolation in the near future. Please specify the interpolation mode explicitly if you plan to keep using nearest interpolation.
warn(
```
## Train the network
To simulate a realistic self-supervised learning scenario, we do not use any supervised metrics for training,
such as PSNR or SSIM, which require clean ground truth images.
#### TIP
We can use the same self-supervised loss for evaluation, as it does not require clean images,
to monitor the training process (e.g. for early stopping). This is done automatically when `metrics=None` and `early_stop>0` in the trainer.
```Python
verbose = True # print training information
train_dataloader = DataLoader(
train_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True
)
test_dataloader = DataLoader(
test_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False
)
# Initialize the trainer
trainer = dinv.Trainer(
model,
physics=physics,
epochs=epochs,
scheduler=scheduler,
losses=losses,
optimizer=optimizer,
train_dataloader=train_dataloader,
eval_dataloader=test_dataloader,
compute_eval_losses=True, # use self-supervised loss for evaluation
early_stop_on_losses=True, # stop using self-supervised eval loss
metrics=None, # no supervised metrics
early_stop=2, # early stop using the self-supervised loss on the test set
plot_images=True,
device=device,
save_path=str(CKPT_DIR / operation),
verbose=verbose,
show_progress_bar=False, # disable progress bar for better vis in sphinx gallery.
ckp_interval=10,
)
model = trainer.train()
```
```none
/local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1415: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
self.setup_train()
The model has 187019 trainable parameters
Train epoch 0: MCLoss=0.0, EILoss=0.0, TotalLoss=0.0
/local/jtachell/deepinv/deepinv/.pixi/envs/docs/lib/python3.12/multiprocessing/popen_fork.py:66: DeprecationWarning: This process (pid=4088353) is multi-threaded, use of fork() may lead to deadlocks in the child.
self.pid = os.fork()
Eval epoch 0: MCLoss=0.0, EILoss=0.0, TotalLoss=0.0
Best model saved at epoch 1
```
## Test the network
We now assume that we have access to a small test set of ground-truth images to evaluate the performance of the trained network.
and we compute the PSNR between the denoised images and the clean ground truth images.
```Python
trainer.test(test_dataloader, metrics=dinv.metric.PSNR())
```
```none
/local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1611: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
self.setup_train(train=False)
/local/jtachell/deepinv/deepinv/.pixi/envs/docs/lib/python3.12/multiprocessing/popen_fork.py:66: DeprecationWarning: This process (pid=4088353) is multi-threaded, use of fork() may lead to deadlocks in the child.
self.pid = os.fork()
Eval epoch 0: MCLoss=0.0, EILoss=0.0, TotalLoss=0.0, PSNR=37.319, PSNR no learning=32.749
Test results:
PSNR no learning: 32.749 +- 0.000
PSNR: 37.319 +- 0.000
{'PSNR no learning': 32.74856948852539, 'PSNR no learning_std': 0.0, 'PSNR': 37.318851470947266, 'PSNR_std': 0.0, 'runtime': 0.0332610672339797, 'peak_gpu_memory_usage': 1.1032156944274902}
```
* **References:**
* Dongdong Chen, Julián Tachella, and Mike E Davies. Equivariant imaging: learning beyond the range space. In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, 4379–4388. 2021.
* Hemant K Aggarwal, Merry P Mani, and Mathews Jacob. Modl: model-based deep learning architecture for inverse problems. *IEEE transactions on medical imaging*, 38(2):394–405, 2018.
* Julián Tachella, Dongdong Chen, and Mike Davies. Sensing theorems for unsupervised learning in linear inverse problems. *Journal of Machine Learning Research*, 24(39):1–45, 2023. URL: [http://jmlr.org/papers/v24/22-0315.html](http://jmlr.org/papers/v24/22-0315.html).
**Total running time of the script:** (0 minutes 5.963 seconds)
[`Download Jupyter notebook: demo_equivariant_imaging.ipynb`](https://deepinv.org/auto_examples/self-supervised-learning/_downloads/fff43dcb7e01ba08f91d580209308dce/demo_equivariant_imaging.ipynb)
[`Download Python source code: demo_equivariant_imaging.py`](https://deepinv.org/auto_examples/self-supervised-learning/_downloads/cc2dd3ba2cedc98abbfca9b6e4a8dd51/demo_equivariant_imaging.py)
[`Download zipped: demo_equivariant_imaging.zip`](https://deepinv.org/auto_examples/self-supervised-learning/_downloads/b313b423aa792cc8a70bf9aea60f7d63/demo_equivariant_imaging.zip)
[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)