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_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…

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

dataset = dinv.datasets.Urban100HR(
    ".",
    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]
 26%|██▋       | 34.1M/129M [00:00<00:00, 358MB/s]
 56%|█████▌    | 72.1M/129M [00:00<00:00, 381MB/s]
 85%|████████▍ | 110M/129M [00:00<00:00, 387MB/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, 197.97it/s]
Extracting:  47%|████▋     | 47/101 [00:00<00:00, 231.60it/s]
Extracting:  70%|███████   | 71/101 [00:00<00:00, 222.81it/s]
Extracting:  93%|█████████▎| 94/101 [00:00<00:00, 217.70it/s]
Extracting: 100%|██████████| 101/101 [00:00<00:00, 218.14it/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.033, PSNR=15.701
Eval epoch 0: PSNR=20.925
Best model saved at epoch 1
Train epoch 1: TotalLoss=0.005, PSNR=24.101
Eval epoch 1: PSNR=25.376
Best model saved at epoch 2
Train epoch 2: TotalLoss=0.002, PSNR=27.276
Eval epoch 2: PSNR=27.283
Best model saved at epoch 3
Train epoch 3: TotalLoss=0.002, PSNR=27.658
Eval epoch 3: PSNR=27.91
Best model saved at epoch 4
Train epoch 4: TotalLoss=0.002, PSNR=28.647
Eval epoch 4: PSNR=28.848
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=28.848, PSNR no learning=13.9
Test results:
PSNR no learning: 13.900 +- 1.983
PSNR: 28.848 +- 2.042

{'PSNR no learning': 13.900366020202636, 'PSNR no learning_std': 1.9832890603307352, 'PSNR': 28.84770851135254, 'PSNR_std': 2.041819680795037}