.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/self-supervised-learning/demo_unsure.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note New to DeepInverse? Get started with the basics with the :ref:`5 minute quickstart tutorial `.. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_self-supervised-learning_demo_unsure.py: Self-supervised denoising with the UNSURE loss. ==================================================================================================== This example shows you how to train a denoiser network in a fully self-supervised way, i.e., using noisy images with unknown noise level only via the UNSURE loss, which is introduced by :footcite:t:`tachella2024unsure`. The UNSURE optimization problem for Gaussian denoising with unknown noise level is defined as: .. math:: \min_{R} \max_{\sigma^2} \frac{1}{m}\|y-\inverse{y}\|_2^2 +\frac{2\sigma^2}{m\tau}b^{\top} \left(\inverse{y+\tau b}-\inverse{y}\right) where :math:`R` is the trainable network, :math:`y` is the noisy image with :math:`m` pixels, :math:`b\sim \mathcal{N}(0,1)` is a Gaussian random variable, :math:`\tau` is a small positive number, and :math:`\odot` is an elementwise multiplication. .. GENERATED FROM PYTHON SOURCE LINES 19-29 .. code-block:: Python from pathlib import Path import torch from torch.utils.data import DataLoader from torchvision import transforms, datasets import deepinv as dinv from deepinv.utils import get_data_home .. GENERATED FROM PYTHON SOURCE LINES 30-33 Setup paths for data loading and results. --------------------------------------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 33-44 .. code-block:: Python BASE_DIR = Path(".") DATA_DIR = BASE_DIR / "measurements" CKPT_DIR = BASE_DIR / "ckpts" ORIGINAL_DATA_DIR = get_data_home() # Set the global random seed from pytorch to ensure reproducibility of the example. torch.manual_seed(0) device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu" .. rst-class:: sphx-glr-script-out .. code-block:: none Selected GPU 0 with 5939.25 MiB free memory .. GENERATED FROM PYTHON SOURCE LINES 45-49 Load base image datasets ---------------------------------------------------------------------------------- In this example, we use the MNIST dataset as the base image dataset. .. GENERATED FROM PYTHON SOURCE LINES 49-62 .. code-block:: Python operation = "denoising" train_dataset_name = "MNIST" transform = transforms.Compose([transforms.ToTensor()]) train_dataset = datasets.MNIST( root=ORIGINAL_DATA_DIR, train=True, transform=transform, download=True ) test_dataset = datasets.MNIST( root=ORIGINAL_DATA_DIR, train=False, transform=transform, download=True ) .. GENERATED FROM PYTHON SOURCE LINES 63-72 Generate a dataset of noisy images ---------------------------------------------------------------------------------- We generate a dataset of noisy images corrupted by Gaussian noise. .. note:: We use a subset of the whole training set to reduce the computational load of the example. We recommend to use the whole set by setting ``n_images_max=None`` to get the best results. .. GENERATED FROM PYTHON SOURCE LINES 72-102 .. code-block:: Python true_sigma = 0.1 # defined physics physics = dinv.physics.Denoising(dinv.physics.GaussianNoise(sigma=true_sigma)) # 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 n_images_max = ( 100 if torch.cuda.is_available() else 5 ) # number of images used for training measurement_dir = DATA_DIR / train_dataset_name / 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, test_datapoints=n_images_max, num_workers=num_workers, dataset_filename="demo_sure", ) train_dataset = dinv.datasets.HDF5Dataset(path=deepinv_datasets_path, train=True) test_dataset = dinv.datasets.HDF5Dataset(path=deepinv_datasets_path, train=False) .. rst-class:: sphx-glr-script-out .. code-block:: none /local/jtachell/deepinv/deepinv/deepinv/datasets/datagenerator.py:344: UserWarning: Dataset measurements/MNIST/denoising/demo_sure0.h5 already exists, this will close and overwrite the previous dataset. warn( Dataset has been saved at measurements/MNIST/denoising/demo_sure0.h5 .. GENERATED FROM PYTHON SOURCE LINES 103-107 Set up the denoiser network --------------------------------------------------------------- We use a simple U-Net architecture with 2 scales as the denoiser network. .. GENERATED FROM PYTHON SOURCE LINES 107-113 .. code-block:: Python model = dinv.models.ArtifactRemoval( dinv.models.UNet(in_channels=1, out_channels=1, scales=2).to(device) ) .. GENERATED FROM PYTHON SOURCE LINES 114-128 Set up the training parameters -------------------------------------------- We set :class:`deepinv.loss.SureGaussianLoss` as the training loss with the ``unsure=True`` option. The optimization with respect to the noise level is done by stochastic gradient descent with momentum inside the loss class, so it is seamlessly integrated into the training process. .. note:: There are (UN)SURE losses for various noise distributions. See also :class:`deepinv.loss.SurePGLoss` for mixed Poisson-Gaussian noise. .. note:: We train for only 10 epochs to reduce the computational load of the example. We recommend to train for more epochs to get the best results. .. GENERATED FROM PYTHON SOURCE LINES 128-147 .. code-block:: Python epochs = 10 # choose training epochs learning_rate = 5e-4 batch_size = 32 if torch.cuda.is_available() else 1 sigma_init = 0.05 # initial guess for the noise level step_size = 1e-4 # step size for the optimization of the noise level momentum = 0.9 # momentum for the optimization of the noise level # choose self-supervised training loss loss = dinv.loss.SureGaussianLoss( sigma=sigma_init, unsure=True, step_size=step_size, momentum=momentum ) # choose optimizer and scheduler optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-8) print(f"INIT. noise level {loss.sigma2.sqrt().item():.3f}") .. rst-class:: sphx-glr-script-out .. code-block:: none INIT. noise level 0.050 .. GENERATED FROM PYTHON SOURCE LINES 148-160 Train the network -------------------------------------------- We train the network using the :class:`deepinv.Trainer` class. 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 (without updating the noise level, which is equivalent to SURE with the estimated noise level), 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. .. GENERATED FROM PYTHON SOURCE LINES 160-189 .. code-block:: Python train_dataloader = DataLoader( train_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True ) # Initialize the trainer trainer = dinv.Trainer( model=model, physics=physics, epochs=epochs, losses=loss, 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 optimizer=optimizer, device=device, train_dataloader=train_dataloader, eval_dataloader=train_dataloader, plot_images=False, save_path=str(CKPT_DIR / operation), verbose=True, # print training information show_progress_bar=False, # disable progress bar for better vis in sphinx gallery. ) # Train the network model = trainer.train() .. rst-class:: sphx-glr-script-out .. code-block:: none /local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1337: 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 444737 trainable parameters Train epoch 0: TotalLoss=0.237 Eval epoch 0: TotalLoss=0.086 Best model saved at epoch 1 Train epoch 1: TotalLoss=0.061 Eval epoch 1: TotalLoss=0.032 Best model saved at epoch 2 Train epoch 2: TotalLoss=0.023 Eval epoch 2: TotalLoss=0.016 Best model saved at epoch 3 Train epoch 3: TotalLoss=0.013 Eval epoch 3: TotalLoss=0.011 Best model saved at epoch 4 Train epoch 4: TotalLoss=0.01 Eval epoch 4: TotalLoss=0.009 Best model saved at epoch 5 Train epoch 5: TotalLoss=0.01 Eval epoch 5: TotalLoss=0.009 Best model saved at epoch 6 Train epoch 6: TotalLoss=0.008 Eval epoch 6: TotalLoss=0.007 Best model saved at epoch 7 Train epoch 7: TotalLoss=0.006 Eval epoch 7: TotalLoss=0.006 Best model saved at epoch 8 Train epoch 8: TotalLoss=0.005 Eval epoch 8: TotalLoss=0.004 Best model saved at epoch 9 Train epoch 9: TotalLoss=0.004 Eval epoch 9: TotalLoss=0.003 Best model saved at epoch 10 .. GENERATED FROM PYTHON SOURCE LINES 190-194 Check learned noise level -------------------------------------------- We can verify the learned noise level by checking the estimated noise level from the loss function. .. GENERATED FROM PYTHON SOURCE LINES 194-201 .. code-block:: Python est_sigma = loss.sigma2.sqrt().item() print(f"LEARNED noise level {est_sigma:.3f}") print(f"Estimation error noise level {abs(est_sigma-true_sigma):.3f}") .. rst-class:: sphx-glr-script-out .. code-block:: none LEARNED noise level 0.115 Estimation error noise level 0.015 .. GENERATED FROM PYTHON SOURCE LINES 202-205 Test the network -------------------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 205-213 .. code-block:: Python test_dataloader = DataLoader( test_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False ) trainer.plot_images = True trainer.test(test_dataloader=test_dataloader, metrics=dinv.metric.PSNR()) .. image-sg:: /auto_examples/self-supervised-learning/images/sphx_glr_demo_unsure_001.png :alt: Ground truth, Measurement, No learning, Reconstruction :srcset: /auto_examples/self-supervised-learning/images/sphx_glr_demo_unsure_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none /local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1529: 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: TotalLoss=0.003, PSNR=23.169, PSNR no learning=20.059 Test results: PSNR no learning: 20.059 +- 0.224 PSNR: 23.169 +- 0.629 {'PSNR no learning': 20.058787994384765, 'PSNR no learning_std': 0.2240344452393849, 'PSNR': 23.168941574096678, 'PSNR_std': 0.6285484021991131} .. GENERATED FROM PYTHON SOURCE LINES 214-217 :References: .. footbibliography:: .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 11.052 seconds) .. _sphx_glr_download_auto_examples_self-supervised-learning_demo_unsure.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: demo_unsure.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: demo_unsure.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: demo_unsure.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_