#### 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).. # 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](https://deepinv.org/user_guide/reconstruction/introduction.html.md#reconstructors) or define your own. * A `physics` from our [list of physics](https://deepinv.org/user_guide/physics/physics.html.md#physics). Or, [bring your own physics](https://deepinv.org/auto_examples/basics/demo_custom_dataset.html.md#sphx-glr-auto-examples-basics-demo-custom-dataset-py). * A `dataset` of images and/or measurements from [datasets](https://deepinv.org/user_guide/training/datasets.html.md#datasets). Or, [bring your own dataset](https://deepinv.org/auto_examples/basics/demo_custom_dataset.html.md#sphx-glr-auto-examples-basics-demo-custom-dataset-py). * A `loss` from our [loss functions](https://deepinv.org/user_guide/training/loss.html.md#loss). * A `metric` from our [metrics](https://deepinv.org/user_guide/training/metric.html.md#metric). Here, we demonstrate a simple experiment of training a UNet on an inpainting task on the Urban100 dataset of natural images. ```Python import deepinv as dinv import torch device = dinv.utils.get_device() rng = torch.Generator(device=device).manual_seed(0) ``` ```none Selected GPU 0 with 28660.0 MiB free memory ``` ## Setup First, define the physics that we want to train on. ```Python 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](https://deepinv.org/user_guide/training/datasets.html.md#datasets) for types of datasets DeepInverse supports: e.g. paired, ground-truth-free, single-image… ```Python 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, use_dict_output=True), shuffle=True, ) test_dataloader = torch.utils.data.DataLoader( dinv.datasets.HDF5Dataset(dataset_path, train=False, use_dict_output=True), shuffle=False, ) ``` ```none Dataset has been saved at ./dinv_dataset0.h5 ``` Visualize a data sample: ```Python batch = next(iter(test_dataloader)) x, y = batch["x"], batch["y"] 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)) $$ ```Python 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`](https://deepinv.org/api/stubs/deepinv.Trainer.html.md#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](https://deepinv.org/user_guide/training/loss.html.md#loss) for all supported state-of-the-art loss functions. We evaluate using the PSNR metric. See [metrics](https://deepinv.org/user_guide/training/metric.html.md#metric) 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. ```Python 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() ``` ```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 443585 trainable parameters Train epoch 0: TotalLoss=0.043, PSNR=14.156 Eval epoch 0: PSNR=15.734 Best model saved at epoch 1 Train epoch 1: TotalLoss=0.008, PSNR=21.557 Eval epoch 1: PSNR=23.577 Best model saved at epoch 2 Train epoch 2: TotalLoss=0.003, PSNR=26.129 Eval epoch 2: PSNR=26.903 Best model saved at epoch 3 Train epoch 3: TotalLoss=0.002, PSNR=28.031 Eval epoch 3: PSNR=27.329 Best model saved at epoch 4 Train epoch 4: TotalLoss=0.001, PSNR=29.054 Eval epoch 4: PSNR=28.805 Best model saved at epoch 5 ``` ## Test the network We can now test the trained network using the [`deepinv.test()`](https://deepinv.org/api/stubs/deepinv.test.html.md#deepinv.test) function. The testing function will compute metrics and plot and save the results. ```Python trainer.test(test_dataloader) ``` ```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) Eval epoch 0: PSNR=28.805, PSNR no learning=12.644 Test results: PSNR no learning: 12.644 +- 1.565 PSNR: 28.805 +- 2.564 {'PSNR no learning': 12.64423770904541, 'PSNR no learning_std': 1.5653784165460078, 'PSNR': 28.805314445495604, 'PSNR_std': 2.5636422909079197, 'runtime': 0.0270218919031322, 'peak_gpu_memory_usage': 0.045529842376708984} ``` **Total running time of the script:** (0 minutes 14.074 seconds) [`Download Jupyter notebook: demo_training.ipynb`](https://deepinv.org/auto_examples/models/_downloads/e93bd4a5161836ee3c97f5f4cb481367/demo_training.ipynb) [`Download Python source code: demo_training.py`](https://deepinv.org/auto_examples/models/_downloads/d8bf7307640fa4af78683a7f34990d0a/demo_training.py) [`Download zipped: demo_training.zip`](https://deepinv.org/auto_examples/models/_downloads/4a9348801bbd60ebc60c98088a1cae05/demo_training.zip) [Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)