.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/unfolded/demo_vanilla_unfolded.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 <sphx_glr_auto_examples_basics_demo_quickstart.py>`..

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

.. _sphx_glr_auto_examples_unfolded_demo_vanilla_unfolded.py:


Vanilla Unfolded algorithm for super-resolution
====================================================================================================

This is a simple example to show how to use vanilla unfolded Plug-and-Play.
The DnCNN denoiser and the algorithm parameters (stepsize, regularization parameters) are trained jointly.
For simplicity, we show how to train the algorithm on a  small dataset. For optimal results, use a larger dataset.

.. GENERATED FROM PYTHON SOURCE LINES 11-22

.. code-block:: Python

    import deepinv as dinv
    import torch
    from deepinv.models.utils import get_weights_url
    from torch.utils.data import DataLoader
    from deepinv.optim.data_fidelity import L2
    from deepinv.optim.prior import PnP
    from deepinv.optim import DRS
    from torchvision import transforms
    from deepinv.utils import get_data_home
    from deepinv.datasets import BSDS500








.. GENERATED FROM PYTHON SOURCE LINES 23-26

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


.. GENERATED FROM PYTHON SOURCE LINES 26-37

.. code-block:: Python


    BASE_DIR = get_data_home()
    DATA_DIR = BASE_DIR / "measurements"
    RESULTS_DIR = BASE_DIR / "results"
    CKPT_DIR = BASE_DIR / "ckpts"

    # 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 5036.125 MiB free memory




.. GENERATED FROM PYTHON SOURCE LINES 38-41

Load base image datasets and degradation operators.
----------------------------------------------------------------------------------------
In this example, we use the CBSD500 dataset for training and the Set3C dataset for testing.

.. GENERATED FROM PYTHON SOURCE LINES 41-46

.. code-block:: Python


    img_size = 64 if torch.cuda.is_available() else 32
    n_channels = 3  # 3 for color images, 1 for gray-scale images
    operation = "super-resolution"








.. GENERATED FROM PYTHON SOURCE LINES 47-50

Generate a dataset of low resolution images and load it.
----------------------------------------------------------------------------------------
We use the Downsampling class from the physics module to generate a dataset of low resolution images.

.. GENERATED FROM PYTHON SOURCE LINES 50-104

.. code-block:: Python


    # For simplicity, we use a small dataset for training.
    # To be replaced for optimal results. For example, you can use the larger DIV2K or LSDIR datasets (also provided in the library).

    # Specify the  train and test transforms to be applied to the input images.
    test_transform = transforms.Compose(
        [transforms.CenterCrop(img_size), transforms.ToTensor()]
    )
    train_transform = transforms.Compose(
        [transforms.RandomCrop(img_size), transforms.ToTensor()]
    )
    # Define the base train and test datasets of clean images.
    train_base_dataset = BSDS500(
        BASE_DIR, download=True, train=True, transform=train_transform
    )
    test_base_dataset = BSDS500(
        BASE_DIR, download=False, train=False, transform=test_transform
    )

    # Use parallel dataloader if using a GPU to speed up training, otherwise, as all computes are on CPU, use synchronous
    # dataloading.
    num_workers = 4 if torch.cuda.is_available() else 0

    # Degradation parameters
    factor = 2
    noise_level_img = 0.03

    # Generate the gaussian blur downsampling operator.
    physics = dinv.physics.Downsampling(
        filter="gaussian",
        img_size=(n_channels, img_size, img_size),
        factor=factor,
        device=device,
        noise_model=dinv.physics.GaussianNoise(sigma=noise_level_img),
    )
    my_dataset_name = "demo_unfolded_sr"
    n_images_max = (
        None if torch.cuda.is_available() else 10
    )  # max number of images used for training (use all if you have a GPU)
    measurement_dir = DATA_DIR / "BSDS500" / operation
    generated_datasets_path = dinv.datasets.generate_dataset(
        train_dataset=train_base_dataset,
        test_dataset=test_base_dataset,
        physics=physics,
        device=device,
        save_dir=measurement_dir,
        train_datapoints=n_images_max,
        num_workers=num_workers,
        dataset_filename=str(my_dataset_name),
    )

    train_dataset = dinv.datasets.HDF5Dataset(path=generated_datasets_path, train=True)
    test_dataset = dinv.datasets.HDF5Dataset(path=generated_datasets_path, train=False)





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

 .. code-block:: none


    0it [00:00, ?it/s]
    1.06MB [00:00, 11.1MB/s]
    5.62MB [00:00, 32.5MB/s]
    8.75MB [00:00, 28.7MB/s]
    11.8MB [00:00, 29.6MB/s]
    14.6MB [00:00, 29.4MB/s]
    17.5MB [00:00, 28.9MB/s]
    20.3MB [00:00, 28.7MB/s]
    23.5MB [00:00, 29.5MB/s]
    26.9MB [00:00, 30.6MB/s]
    30.4MB [00:01, 32.4MB/s]
    33.8MB [00:01, 33.1MB/s]
    36.9MB [00:01, 32.2MB/s]
    40.1MB [00:01, 32.2MB/s]
    43.6MB [00:01, 33.3MB/s]
    46.8MB [00:01, 33.4MB/s]
    50.2MB [00:01, 33.9MB/s]
    53.4MB [00:01, 33.8MB/s]
    56.7MB [00:01, 31.3MB/s]
    59.8MB [00:02, 25.5MB/s]
    62.4MB [00:02, 23.0MB/s]
    64.8MB [00:02, 21.3MB/s]
    66.9MB [00:02, 20.1MB/s]
    68.9MB [00:02, 19.8MB/s]
    70.9MB [00:02, 19.2MB/s]
    72.8MB [00:02, 18.1MB/s]
    74.5MB [00:02, 17.3MB/s]
    76.2MB [00:03, 16.3MB/s]
    77.8MB [00:03, 16.3MB/s]
    79.5MB [00:03, 16.6MB/s]
    81.2MB [00:03, 17.1MB/s]
    82.9MB [00:03, 16.7MB/s]
    84.7MB [00:03, 17.2MB/s]
    86.4MB [00:03, 17.1MB/s]
    88.1MB [00:03, 16.9MB/s]
    89.9MB [00:03, 17.5MB/s]
    91.6MB [00:04, 17.5MB/s]
    93.4MB [00:04, 17.7MB/s]
    95.1MB [00:04, 17.8MB/s]
    96.9MB [00:04, 16.8MB/s]
    98.8MB [00:04, 17.3MB/s]
    101MB [00:04, 17.5MB/s] 
    102MB [00:04, 17.4MB/s]
    104MB [00:04, 17.3MB/s]
    106MB [00:04, 17.3MB/s]
    108MB [00:05, 17.6MB/s]
    109MB [00:05, 17.9MB/s]
    111MB [00:05, 18.0MB/s]
    113MB [00:05, 17.9MB/s]
    115MB [00:05, 17.4MB/s]
    116MB [00:05, 17.6MB/s]
    118MB [00:05, 17.1MB/s]
    120MB [00:05, 17.8MB/s]
    122MB [00:05, 17.8MB/s]
    124MB [00:05, 17.8MB/s]
    125MB [00:06, 17.8MB/s]
    127MB [00:06, 17.6MB/s]
    129MB [00:06, 17.5MB/s]
    131MB [00:06, 17.9MB/s]
    133MB [00:06, 18.8MB/s]
    134MB [00:06, 9.61MB/s]
    138MB [00:07, 15.0MB/s]
    140MB [00:07, 15.5MB/s]
    142MB [00:07, 14.5MB/s]
    144MB [00:07, 14.4MB/s]
    146MB [00:07, 15.0MB/s]
    148MB [00:07, 15.4MB/s]
    149MB [00:07, 16.5MB/s]
    152MB [00:07, 17.8MB/s]
    154MB [00:07, 19.4MB/s]
    156MB [00:08, 20.9MB/s]
    158MB [00:08, 16.2MB/s]
    160MB [00:08, 15.5MB/s]
    160MB [00:08, 20.0MB/s]

    Extracting:   0%|          | 0/2492 [00:00<?, ?it/s]
    Extracting:   4%|▍         | 94/2492 [00:00<00:02, 931.66it/s]
    Extracting:   8%|▊         | 190/2492 [00:00<00:02, 945.50it/s]
    Extracting:  12%|█▏        | 289/2492 [00:00<00:02, 961.67it/s]
    Extracting:  15%|█▌        | 386/2492 [00:00<00:02, 914.90it/s]
    Extracting:  19%|█▉        | 484/2492 [00:00<00:02, 935.34it/s]
    Extracting:  23%|██▎       | 582/2492 [00:00<00:02, 948.06it/s]
    Extracting:  28%|██▊       | 698/2492 [00:00<00:01, 1013.23it/s]
    Extracting:  34%|███▎      | 838/2492 [00:00<00:01, 1108.15it/s]
    Extracting:  40%|████      | 1004/2492 [00:00<00:01, 1275.11it/s]
    Extracting:  45%|████▌     | 1132/2492 [00:01<00:01, 746.01it/s] 
    Extracting:  49%|████▉     | 1233/2492 [00:01<00:02, 594.15it/s]
    Extracting:  53%|█████▎    | 1314/2492 [00:01<00:02, 503.87it/s]
    Extracting:  55%|█████▌    | 1381/2492 [00:01<00:02, 457.04it/s]
    Extracting:  58%|█████▊    | 1438/2492 [00:02<00:02, 430.39it/s]
    Extracting:  60%|█████▉    | 1489/2492 [00:02<00:02, 408.58it/s]
    Extracting:  62%|██████▏   | 1535/2492 [00:02<00:02, 404.17it/s]
    Extracting:  66%|██████▌   | 1639/2492 [00:02<00:01, 529.79it/s]
    Extracting:  68%|██████▊   | 1699/2492 [00:02<00:01, 404.47it/s]
    Extracting: 100%|██████████| 2492/2492 [00:02<00:00, 874.44it/s]
    Dataset has been saved at datasets/measurements/BSDS500/super-resolution/demo_unfolded_sr0.h5




.. GENERATED FROM PYTHON SOURCE LINES 105-111

Define the unfolded PnP algorithm.
----------------------------------------------------------------------------------------
The chosen algorithm is here DRS (Douglas-Rachford Splitting).
Note that if the prior (resp. a parameter) is initialized with a list of length max_iter,
then a distinct model (resp. parameter) is trained for each iteration.
For fixed trained model prior (resp. parameter) across iterations, initialize with a single element.

.. GENERATED FROM PYTHON SOURCE LINES 111-149

.. code-block:: Python


    # Unrolled optimization algorithm parameters
    max_iter = 5  # number of unfolded layers

    # Select the data fidelity term
    data_fidelity = L2()

    # Set up the trainable denoising prior
    # Here the prior model is common for all iterations
    prior = PnP(denoiser=dinv.models.DnCNN(depth=20, pretrained="download").to(device))

    # The parameters are initialized with a list of length max_iter, so that a distinct parameter is trained for each iteration.
    stepsize = [1.0] * max_iter  # stepsize of the algorithm
    sigma_denoiser = [
        1.0
    ] * max_iter  # noise level parameter of the denoiser (not used by DnCNN)
    beta = 1.0  # relaxation parameter of the Douglas-Rachford splitting
    trainable_params = [
        "stepsize",
        "beta",
        "sigma_denoiser",
    ]  # define which parameters are trainable

    # Logging parameters
    verbose = True

    # Define the unfolded trainable model.
    model = DRS(
        stepsize=stepsize,
        sigma_denoiser=sigma_denoiser,
        beta=beta,
        trainable_params=trainable_params,
        data_fidelity=data_fidelity,
        max_iter=max_iter,
        prior=prior,
        unfold=True,
    )








.. GENERATED FROM PYTHON SOURCE LINES 150-153

Define the training parameters.
----------------------------------------------------------------------------------------
We use the Adam optimizer and the StepLR scheduler.

.. GENERATED FROM PYTHON SOURCE LINES 153-184

.. code-block:: Python



    # training parameters
    epochs = 5 if torch.cuda.is_available() else 1
    learning_rate = 5e-4
    train_batch_size = 32 if torch.cuda.is_available() else 1
    test_batch_size = 3

    # choose optimizer and scheduler
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-8)

    # If working on CPU, start with a pretrained model to reduce training time
    if not torch.cuda.is_available():
        file_name = "demo_vanilla_unfolded.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
        )
        model.load_state_dict(ckpt["state_dict"])
        optimizer.load_state_dict(ckpt["optimizer"])

    # choose supervised training loss
    losses = [dinv.loss.SupLoss(metric=dinv.metric.MSE())]

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








.. GENERATED FROM PYTHON SOURCE LINES 185-188

Train the network
----------------------------------------------------------------------------------------
We train the network using the :class:`deepinv.Trainer` class.

.. GENERATED FROM PYTHON SOURCE LINES 188-207

.. code-block:: Python


    trainer = dinv.Trainer(
        model,
        physics=physics,
        train_dataloader=train_dataloader,
        eval_dataloader=test_dataloader,
        epochs=epochs,
        losses=losses,
        optimizer=optimizer,
        device=device,
        early_stop=True,  # set to None to disable early stopping
        save_path=str(CKPT_DIR / operation),
        verbose=verbose,
        show_progress_bar=False,  # disable progress bar for better vis in sphinx gallery.
    )

    model = trainer.train()






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

 .. code-block:: none

    /local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1352: 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 668238 trainable parameters
    /local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:514: UserWarning: early_stop should be an integer or None. Setting early_stop=3. This behaviour will be deprecated in future versions.
      warnings.warn(
    Train epoch 0: TotalLoss=0.009, PSNR=21.591
    Eval epoch 0: PSNR=20.581
    Best model saved at epoch 1
    Train epoch 1: TotalLoss=0.007, PSNR=23.006
    Eval epoch 1: PSNR=21.237
    Best model saved at epoch 2
    Train epoch 2: TotalLoss=0.007, PSNR=23.412
    Eval epoch 2: PSNR=21.421
    Best model saved at epoch 3
    Train epoch 3: TotalLoss=0.006, PSNR=23.779
    Eval epoch 3: PSNR=21.717
    Best model saved at epoch 4
    Train epoch 4: TotalLoss=0.006, PSNR=24.052
    Eval epoch 4: PSNR=21.854
    Best model saved at epoch 5




.. GENERATED FROM PYTHON SOURCE LINES 208-212

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



.. GENERATED FROM PYTHON SOURCE LINES 212-230

.. code-block:: Python

    trainer.test(test_dataloader)

    test_sample, _ = next(iter(test_dataloader))
    model.eval()
    test_sample = test_sample.to(device)

    # Get the measurements and the ground truth
    y = physics(test_sample)
    with torch.no_grad():
        rec = model(y, physics=physics)

    backprojected = physics.A_adjoint(y)

    dinv.utils.plot(
        [backprojected, rec, test_sample],
        titles=["Linear", "Reconstruction", "Ground truth"],
        suptitle="Reconstruction results",
    )



.. image-sg:: /auto_examples/unfolded/images/sphx_glr_demo_vanilla_unfolded_001.png
   :alt: Reconstruction results, Linear, Reconstruction, Ground truth
   :srcset: /auto_examples/unfolded/images/sphx_glr_demo_vanilla_unfolded_001.png
   :class: sphx-glr-single-img


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

 .. code-block:: none

    /local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1544: 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=21.854, PSNR no learning=9.122
    Test results:
    PSNR no learning: 9.122 +- 2.903
    PSNR: 21.854 +- 3.672
    /local/jtachell/deepinv/deepinv/deepinv/utils/plotting.py:387: UserWarning: This figure was using a layout engine that is incompatible with subplots_adjust and/or tight_layout; not calling subplots_adjust.
      fig.subplots_adjust(top=0.75)





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

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


.. _sphx_glr_download_auto_examples_unfolded_demo_vanilla_unfolded.py:

.. only:: html

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

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

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

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

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

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

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


.. only:: html

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

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