.. 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_n2n_denoising.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_self-supervised-learning_demo_n2n_denoising.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_self-supervised-learning_demo_n2n_denoising.py:


Self-supervised denoising with the Neighbor2Neighbor loss.
====================================================================================================

This example shows you how to train a denoiser network in a fully self-supervised way,
i.e., using noisy images only via the Neighbor2Neighbor loss, which exploits the local correlation of natural images.

The Neighbor2Neighbor loss is presented in `"Neighbor2Neighbor: Self-Supervised Denoising from Single Noisy Images"
<https://openaccess.thecvf.com/content/CVPR2021/papers/Huang_Neighbor2Neighbor_Self-Supervised_Denoising_From_Single_Noisy_Images_CVPR_2021_paper.pdf>`_
and is defined as:

.. math::

    \| A_2 y - R(A_1 y)\|^2 + \gamma \| A_2 y - R(A_1 y) - (A_2 R(y) - A_1 R(y))\|^2

where :math:`A_1` and :math:`A_2` are two masks, each choosing a different neighboring map,
:math:`R` is the trainable denoiser network, :math:`\gamma>0` is a regularization parameter
and no gradient is propagated when computing :math:`R(y)`.

.. GENERATED FROM PYTHON SOURCE LINES 22-33

.. 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.models.utils import get_weights_url
    from deepinv.utils.demo import get_data_home








.. GENERATED FROM PYTHON SOURCE LINES 34-37

Setup paths for data loading and results.
---------------------------------------------------------------


.. GENERATED FROM PYTHON SOURCE LINES 37-48

.. 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"








.. GENERATED FROM PYTHON SOURCE LINES 49-53

Load base image datasets
----------------------------------------------------------------------------------
In this example, we use the MNIST dataset as the base image dataset.


.. GENERATED FROM PYTHON SOURCE LINES 53-66

.. 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 67-76

Generate a dataset of noisy images
----------------------------------------------------------------------------------

We generate a dataset of noisy images corrupted by Poisson 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 76-105

.. code-block:: Python


    # defined physics
    physics = dinv.physics.Denoising(dinv.physics.PoissonNoise(0.1))

    # Use parallel dataloader if using a GPU to fasten 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

    my_dataset_name = "demo_n2n"
    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=str(my_dataset_name),
    )

    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

    Dataset has been saved at measurements/MNIST/denoising/demo_n2n0.h5




.. GENERATED FROM PYTHON SOURCE LINES 106-110

Set up the denoiser network
---------------------------------------------------------------

We use a simple U-Net architecture with 2 scales as the denoiser network.

.. GENERATED FROM PYTHON SOURCE LINES 110-116

.. 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 117-125

Set up the training parameters
--------------------------------------------
We set :class:`deepinv.loss.Neighbor2Neighbor` as the training loss.

.. note::

      We use a pretrained model to reduce training time. You can get the same results by training from scratch
      for 50 epochs.

.. GENERATED FROM PYTHON SOURCE LINES 125-147

.. code-block:: Python


    epochs = 1  # choose training epochs
    learning_rate = 5e-4
    batch_size = 32 if torch.cuda.is_available() else 1

    # choose self-supervised training loss
    loss = dinv.loss.Neighbor2Neighbor()

    # choose optimizer and scheduler
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-8)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(epochs * 0.8) + 1)

    # start with a pretrained model to reduce training time
    file_name = "ckp_50_demo_n2n.pth"
    url = get_weights_url(model_name="demo", file_name=file_name)
    ckpt = torch.hub.load_state_dict_from_url(
        url, map_location=lambda storage, loc: storage, file_name=file_name
    )
    # load a checkpoint to reduce training time
    model.load_state_dict(ckpt["state_dict"])
    optimizer.load_state_dict(ckpt["optimizer"])





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Downloading: "https://huggingface.co/deepinv/demo/resolve/main/ckp_50_demo_n2n.pth?download=true" to /home/runner/.cache/torch/hub/checkpoints/ckp_50_demo_n2n.pth

      0%|          | 0.00/5.14M [00:00<?, ?B/s]
     44%|████▍     | 2.25M/5.14M [00:00<00:00, 21.7MB/s]
    100%|█████████▉| 5.12M/5.14M [00:00<00:00, 25.2MB/s]
    100%|██████████| 5.14M/5.14M [00:00<00:00, 24.5MB/s]




.. GENERATED FROM PYTHON SOURCE LINES 148-152

Train the network
--------------------------------------------



.. GENERATED FROM PYTHON SOURCE LINES 152-183

.. code-block:: Python


    verbose = True  # print training information
    wandb_vis = False  # plot curves and images in Weight&Bias

    train_dataloader = DataLoader(
        train_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True
    )
    test_dataloader = DataLoader(
        test_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False
    )

    # Initialize the trainer
    trainer = dinv.Trainer(
        model=model,
        physics=physics,
        epochs=epochs,
        scheduler=scheduler,
        losses=loss,
        optimizer=optimizer,
        device=device,
        train_dataloader=train_dataloader,
        eval_dataloader=test_dataloader,
        plot_images=True,
        save_path=str(CKPT_DIR / operation),
        verbose=verbose,
        show_progress_bar=False,  # disable progress bar for better vis in sphinx gallery.
        wandb_vis=wandb_vis,
    )

    model = trainer.train()




.. rst-class:: sphx-glr-horizontal


    *

      .. image-sg:: /auto_examples/self-supervised-learning/images/sphx_glr_demo_n2n_denoising_001.png
         :alt: Ground truth, Measurement, Reconstruction
         :srcset: /auto_examples/self-supervised-learning/images/sphx_glr_demo_n2n_denoising_001.png
         :class: sphx-glr-multi-img

    *

      .. image-sg:: /auto_examples/self-supervised-learning/images/sphx_glr_demo_n2n_denoising_002.png
         :alt: Ground truth, Measurement, Reconstruction
         :srcset: /auto_examples/self-supervised-learning/images/sphx_glr_demo_n2n_denoising_002.png
         :class: sphx-glr-multi-img


.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    The model has 444737 trainable parameters
    Train epoch 0: TotalLoss=0.07, PSNR=24.558
    Eval epoch 0: PSNR=24.633




.. GENERATED FROM PYTHON SOURCE LINES 184-188

Test the network
--------------------------------------------



.. GENERATED FROM PYTHON SOURCE LINES 188-190

.. code-block:: Python


    trainer.test(test_dataloader)



.. image-sg:: /auto_examples/self-supervised-learning/images/sphx_glr_demo_n2n_denoising_003.png
   :alt: Ground truth, Measurement, No learning, Reconstruction
   :srcset: /auto_examples/self-supervised-learning/images/sphx_glr_demo_n2n_denoising_003.png
   :class: sphx-glr-single-img


.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Eval epoch 0: PSNR=24.633, PSNR no learning=19.346
    Test results:
    PSNR no learning: 19.346 +- 1.786
    PSNR: 24.633 +- 1.501

    {'PSNR no learning': np.float64(19.346023559570312), 'PSNR no learning_std': np.float64(1.78632784200239), 'PSNR': np.float64(24.633206176757813), 'PSNR_std': np.float64(1.5012940824804186)}




.. rst-class:: sphx-glr-timing

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


.. _sphx_glr_download_auto_examples_self-supervised-learning_demo_n2n_denoising.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: demo_n2n_denoising.ipynb <demo_n2n_denoising.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: demo_n2n_denoising.py <demo_n2n_denoising.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: demo_n2n_denoising.zip <demo_n2n_denoising.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_