Training a reconstruction model#

This example provides a very simple quick start introduction to training reconstruction networks with DeepInverse for solving imaging inverse problems.

Training requires these components, all of which you can define with DeepInverse:

Here, we demonstrate a simple experiment of training a UNet on an inpainting task on the Urban100 dataset of natural images.

import deepinv as dinv
import torch

device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"
rng = torch.Generator(device=device).manual_seed(0)

Setup#

First, define the physics that we want to train on.

physics = dinv.physics.Inpainting((1, 64, 64), mask=0.8, device=device, rng=rng)

Then define the dataset. Here we simulate a dataset of measurements from Urban100.

Tip

See datasets for types of datasets DeepInverse supports: e.g. paired, ground-truth-free, single-image…

  0%|          | 0/135388067 [00:00<?, ?it/s]
 27%|██▋       | 35.5M/129M [00:00<00:00, 372MB/s]
 56%|█████▌    | 72.4M/129M [00:00<00:00, 381MB/s]
 86%|████████▌ | 111M/129M [00:00<00:00, 392MB/s]
100%|██████████| 129M/129M [00:00<00:00, 383MB/s]

Extracting:   0%|          | 0/101 [00:00<?, ?it/s]
Extracting:  21%|██        | 21/101 [00:00<00:00, 200.75it/s]
Extracting:  47%|████▋     | 47/101 [00:00<00:00, 233.95it/s]
Extracting:  70%|███████   | 71/101 [00:00<00:00, 224.93it/s]
Extracting:  93%|█████████▎| 94/101 [00:00<00:00, 219.16it/s]
Extracting: 100%|██████████| 101/101 [00:00<00:00, 219.75it/s]
Dataset has been successfully downloaded.
Dataset has been saved at ./dinv_dataset0.h5

Visualize a data sample:

x, y = next(iter(test_dataloader))
dinv.utils.plot({"Ground truth": x, "Measurement": y, "Mask": physics.mask})
Ground truth, Measurement, Mask

For the model we use an artifact removal model, where \(\phi_{\theta}\) is a U-Net

\[f_{\theta}(y) = \phi_{\theta}(A^{\top}(y))\]
model = dinv.models.ArtifactRemoval(
    dinv.models.UNet(1, 1, scales=2, batch_norm=False).to(device)
)

Train the model#

We train the model using the deepinv.Trainer class, which cleanly handles all steps for training.

We perform supervised learning and use the mean squared error as loss function. See losses for all supported state-of-the-art loss functions.

We evaluate using the PSNR metric. See metrics for all supported metrics.

Note

In this example, we only train for a few epochs to keep the training time short. For a good reconstruction quality, we recommend to train for at least 100 epochs.

trainer = dinv.Trainer(
    model=model,
    physics=physics,
    optimizer=torch.optim.Adam(model.parameters(), lr=1e-3),
    train_dataloader=train_dataloader,
    eval_dataloader=test_dataloader,
    epochs=5,
    losses=dinv.loss.SupLoss(metric=dinv.metric.MSE()),
    metrics=dinv.metric.PSNR(),
    device=device,
    plot_images=True,
    show_progress_bar=False,
)

_ = trainer.train()
  • Ground truth, Measurement, Reconstruction
  • Ground truth, Measurement, Reconstruction
  • Ground truth, Measurement, Reconstruction
  • Ground truth, Measurement, Reconstruction
  • Ground truth, Measurement, Reconstruction
  • Ground truth, Measurement, Reconstruction
  • Ground truth, Measurement, Reconstruction
  • Ground truth, Measurement, Reconstruction
  • Ground truth, Measurement, Reconstruction
  • Ground truth, Measurement, Reconstruction
The model has 443585 trainable parameters
Train epoch 0: TotalLoss=0.018, PSNR=19.531
Eval epoch 0: PSNR=24.2
Best model saved at epoch 1
Train epoch 1: TotalLoss=0.004, PSNR=25.233
Eval epoch 1: PSNR=26.576
Best model saved at epoch 2
Train epoch 2: TotalLoss=0.002, PSNR=27.892
Eval epoch 2: PSNR=24.744
Train epoch 3: TotalLoss=0.002, PSNR=27.936
Eval epoch 3: PSNR=28.604
Best model saved at epoch 4
Train epoch 4: TotalLoss=0.001, PSNR=29.618
Eval epoch 4: PSNR=30.549
Best model saved at epoch 5

Test the network#

We can now test the trained network using the deepinv.test() function.

The testing function will compute metrics and plot and save the results.

Ground truth, Measurement, No learning, Reconstruction
Eval epoch 0: PSNR=30.549, PSNR no learning=14.167
Test results:
PSNR no learning: 14.167 +- 2.535
PSNR: 30.549 +- 2.943

{'PSNR no learning': np.float64(14.167121887207031), 'PSNR no learning_std': np.float64(2.534591007173441), 'PSNR': np.float64(30.548965454101562), 'PSNR_std': np.float64(2.942676077418184)}

Total running time of the script: (0 minutes 17.982 seconds)

Gallery generated by Sphinx-Gallery