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:

• A model to be trained from reconstructors or define your own.

• A physics from our list of physics. Or, bring your own physics.

• A dataset of images and/or measurements from datasets. Or, bring your own dataset.

• A loss from our loss functions.

• A metric from our metrics.

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_device()
rng = torch.Generator(device=device).manual_seed(0)

Selected CPU device

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…

from torchvision.transforms import Compose, ToTensor, Resize, CenterCrop, Grayscale

dataset = dinv.datasets.Urban100HR(
    dinv.utils.get_cache_home() / "datasets" / "Urban100",
    download=True,
    transform=Compose([ToTensor(), Grayscale(), Resize(256), CenterCrop(64)]),
)

train_dataset, test_dataset = torch.utils.data.random_split(
    torch.utils.data.Subset(dataset, range(50)), (0.8, 0.2)
)

dataset_path = dinv.datasets.generate_dataset(
    train_dataset=train_dataset,
    test_dataset=test_dataset,
    physics=physics,
    device=device,
    save_dir=".",
    batch_size=1,
)

train_dataloader = torch.utils.data.DataLoader(
    dinv.datasets.HDF5Dataset(dataset_path, train=True), shuffle=True
)
test_dataloader = torch.utils.data.DataLoader(
    dinv.datasets.HDF5Dataset(dataset_path, train=False), shuffle=False
)

  0%|          | 0/135388067 [00:00<?, ?it/s]
  7%|▋         | 9.44M/129M [00:00<00:01, 98.6MB/s]
 16%|█▌        | 20.3M/129M [00:00<00:01, 108MB/s]
 24%|██▍       | 31.2M/129M [00:00<00:00, 111MB/s]
 33%|███▎      | 42.1M/129M [00:00<00:00, 112MB/s]
 41%|████      | 53.0M/129M [00:00<00:00, 113MB/s]
 50%|████▉     | 63.9M/129M [00:00<00:00, 113MB/s]
 58%|█████▊    | 74.8M/129M [00:00<00:00, 113MB/s]
 66%|██████▋   | 85.8M/129M [00:00<00:00, 113MB/s]
 75%|███████▍  | 96.6M/129M [00:00<00:00, 113MB/s]
 83%|████████▎ | 108M/129M [00:01<00:00, 113MB/s]
 92%|█████████▏| 118M/129M [00:01<00:00, 114MB/s]
100%|██████████| 129M/129M [00:01<00:00, 112MB/s]

Extracting:   0%|          | 0/101 [00:00<?, ?it/s]
Extracting:  21%|██        | 21/101 [00:00<00:00, 204.30it/s]
Extracting:  47%|████▋     | 47/101 [00:00<00:00, 235.45it/s]
Extracting:  70%|███████   | 71/101 [00:00<00:00, 225.44it/s]
Extracting:  93%|█████████▎| 94/101 [00:00<00:00, 219.41it/s]
Extracting: 100%|██████████| 101/101 [00:00<00:00, 220.59it/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})

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()

• 
• 
• 
• 
• 
• 
• 
• 
• 
• 

The model has 443585 trainable parameters
Train epoch 0: TotalLoss=0.02, PSNR=19.078
Eval epoch 0: PSNR=22.203
Best model saved at epoch 1
Train epoch 1: TotalLoss=0.003, PSNR=25.463
Eval epoch 1: PSNR=25.273
Best model saved at epoch 2
Train epoch 2: TotalLoss=0.002, PSNR=27.989
Eval epoch 2: PSNR=28.602
Best model saved at epoch 3
Train epoch 3: TotalLoss=0.001, PSNR=29.577
Eval epoch 3: PSNR=27.702
Train epoch 4: TotalLoss=0.001, PSNR=30.469
Eval epoch 4: PSNR=30.853
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.

trainer.test(test_dataloader)

Eval epoch 0: PSNR=30.853, PSNR no learning=12.132
Test results:
PSNR no learning: 12.132 +- 1.423
PSNR: 30.853 +- 3.534

{'PSNR no learning': 12.13238458633423, 'PSNR no learning_std': 1.422735333387577, 'PSNR': 30.853326416015626, 'PSNR_std': 3.5338410770998925}