.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/adversarial-learning/demo_gan_imaging.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_adversarial-learning_demo_gan_imaging.py>`
        to download the full example code.

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

.. _sphx_glr_auto_examples_adversarial-learning_demo_gan_imaging.py:


Imaging inverse problems with adversarial networks
==================================================

This example shows you how to train various networks using adversarial
training for deblurring problems. We demonstrate running training and
inference using a conditional GAN (i.e. DeblurGAN), CSGM, AmbientGAN and
UAIR implemented in the library, and how to simply train
your own GAN by using :class:`deepinv.training.AdversarialTrainer`. These
examples can also be easily extended to train more complicated GANs such
as CycleGAN.

This example is based on the following papers:

-  Kupyn et al., `DeblurGAN: Blind Motion Deblurring Using Conditional Adversarial Networks <https://openaccess.thecvf.com/content_cvpr_2018/papers/Kupyn_DeblurGAN_Blind_Motion_CVPR_2018_paper.pdf>`_
-  Bora et al., `Compressed Sensing using Generative
   Models <https://arxiv.org/abs/1703.03208>`_ (CSGM)
-  Bora et al., `AmbientGAN: Generative models from lossy
   measurements <https://openreview.net/forum?id=Hy7fDog0b>`_
-  Pajot et al., `Unsupervised Adversarial Image
   Reconstruction <https://openreview.net/forum?id=BJg4Z3RqF7>`_

Adversarial networks are characterised by the addition of an adversarial
loss :math:`\mathcal{L}_\text{adv}` to the standard reconstruction loss:

.. math:: \mathcal{L}_\text{adv}(x,\hat x;D)=\mathbb{E}_{x\sim p_x}\left[q(D(x))\right]+\mathbb{E}_{\hat x\sim p_{\hat x}}\left[q(1-D(\hat x))\right]

where :math:`D(\cdot)` is the discriminator model, :math:`x` is the
reference image, :math:`\hat x` is the estimated reconstruction,
:math:`q(\cdot)` is a quality function (e.g :math:`q(x)=x` for WGAN).
Training alternates between generator :math:`G` and discriminator
:math:`D` in a minimax game. When there are no ground truths (i.e.
unsupervised), this may be defined on the measurements :math:`y`
instead.

.. GENERATED FROM PYTHON SOURCE LINES 37-58

.. code-block:: Python


    from pathlib import Path

    import torch
    from torch.utils.data import DataLoader, random_split
    from torchvision.datasets import ImageFolder
    from torchvision.transforms import Compose, ToTensor, CenterCrop, Resize
    from torchvision.datasets.utils import download_and_extract_archive

    import deepinv as dinv
    from deepinv.loss import adversarial
    from deepinv.utils.demo import get_data_home
    from deepinv.physics.generator import MotionBlurGenerator

    device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"

    BASE_DIR = Path(".")
    DATA_DIR = BASE_DIR / "measurments"
    ORGINAL_DATA_DIR = get_data_home() / "Urban100"









.. GENERATED FROM PYTHON SOURCE LINES 59-66

Generate dataset
~~~~~~~~~~~~~~~~
In this example we use the Urban100 dataset resized to 128x128. We apply random
motion blur physics using
:class:`deepinv.physics.generator.MotionBlurGenerator`, and save the data
using :func:`deepinv.datasets.generate_dataset`.


.. GENERATED FROM PYTHON SOURCE LINES 66-97

.. code-block:: Python


    physics = dinv.physics.Blur(padding="circular", device=device)
    blur_generator = MotionBlurGenerator((11, 11))

    dataset = dinv.datasets.Urban100HR(
        root=ORGINAL_DATA_DIR,
        download=True,
        transform=Compose([ToTensor(), Resize(256), CenterCrop(128)]),
    )

    train_dataset, test_dataset = random_split(dataset, (0.8, 0.2))

    # Generate data pairs x,y offline using a physics generator
    dataset_path = dinv.datasets.generate_dataset(
        train_dataset=train_dataset,
        test_dataset=test_dataset,
        physics=physics,
        physics_generator=blur_generator,
        device=device,
        save_dir=DATA_DIR,
        batch_size=1,
    )

    train_dataloader = DataLoader(
        dinv.datasets.HDF5Dataset(dataset_path, train=True), shuffle=True
    )
    test_dataloader = DataLoader(
        dinv.datasets.HDF5Dataset(dataset_path, train=False), shuffle=False
    )






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

 .. code-block:: none


      0%|          | 0/135388067 [00:00<?, ?it/s]
     11%|█▏        | 14.8M/129M [00:00<00:00, 154MB/s]
     23%|██▎       | 29.5M/129M [00:00<00:00, 146MB/s]
     42%|████▏     | 54.4M/129M [00:00<00:00, 197MB/s]
     60%|██████    | 77.8M/129M [00:00<00:00, 215MB/s]
     78%|███████▊  | 100M/129M [00:00<00:00, 223MB/s] 
     96%|█████████▌| 123M/129M [00:00<00:00, 229MB/s]
    100%|██████████| 129M/129M [00:00<00:00, 212MB/s]

    Extracting:   0%|          | 0/101 [00:00<?, ?it/s]
    Extracting:  19%|█▉        | 19/101 [00:00<00:00, 187.90it/s]
    Extracting:  40%|███▉      | 40/101 [00:00<00:00, 199.13it/s]
    Extracting:  62%|██████▏   | 63/101 [00:00<00:00, 211.29it/s]
    Extracting:  84%|████████▍ | 85/101 [00:00<00:00, 192.79it/s]
    Extracting: 100%|██████████| 101/101 [00:00<00:00, 196.79it/s]
    Dataset has been successfully downloaded.
    Dataset has been saved at measurments/dinv_dataset0.h5




.. GENERATED FROM PYTHON SOURCE LINES 98-108

Define models
~~~~~~~~~~~~~

We first define reconstruction network (i.e conditional generator) and
discriminator network to use for adversarial training. For demonstration
we use a simple U-Net as the reconstruction network and the
discriminator from `PatchGAN <https://arxiv.org/abs/1611.07004>`_, but
these can be replaced with any architecture e.g transformers, unrolled
etc. Further discriminator models are in :ref:`adversarial models <adversarial>`.


.. GENERATED FROM PYTHON SOURCE LINES 108-135

.. code-block:: Python



    def get_models(model=None, D=None, lr_g=1e-4, lr_d=1e-4, device=device):
        if model is None:
            model = dinv.models.UNet(
                in_channels=3,
                out_channels=3,
                scales=2,
                circular_padding=True,
                batch_norm=False,
            ).to(device)

        if D is None:
            D = dinv.models.PatchGANDiscriminator(n_layers=2, batch_norm=False).to(device)

        optimizer = dinv.training.adversarial.AdversarialOptimizer(
            torch.optim.Adam(model.parameters(), lr=lr_g, weight_decay=1e-8),
            torch.optim.Adam(D.parameters(), lr=lr_d, weight_decay=1e-8),
        )
        scheduler = dinv.training.adversarial.AdversarialScheduler(
            torch.optim.lr_scheduler.StepLR(optimizer.G, step_size=5, gamma=0.9),
            torch.optim.lr_scheduler.StepLR(optimizer.D, step_size=5, gamma=0.9),
        )

        return model, D, optimizer, scheduler









.. GENERATED FROM PYTHON SOURCE LINES 136-156

Conditional GAN training
~~~~~~~~~~~~~~~~~~~~~~~~

Conditional GANs (Kupyn et al., `DeblurGAN: Blind Motion Deblurring Using Conditional Adversarial Networks
<https://openaccess.thecvf.com/content_cvpr_2018/papers/Kupyn_DeblurGAN_Blind_Motion_CVPR_2018_paper.pdf>`_)
are a type of GAN where the generator is conditioned on a label or input. In the context of imaging,
this can be used to generate images from a given measurement. In this example, we use a simple U-Net as the generator
and a PatchGAN discriminator. The forward pass of the generator is given by:

**Conditional GAN** forward pass:

.. math:: \hat x = G(y)

**Conditional GAN** loss:

.. math:: \mathcal{L}=\mathcal{L}_\text{sup}(\hat x, x)+\mathcal{L}_\text{adv}(\hat x, x;D)

where :math:`\mathcal{L}_\text{sup}` is a supervised loss such as
pixel-wise MSE or VGG Perceptual Loss.


.. GENERATED FROM PYTHON SOURCE LINES 156-160

.. code-block:: Python


    G, D, optimizer, scheduler = get_models()









.. GENERATED FROM PYTHON SOURCE LINES 161-165

We next define pixel-wise and adversarial losses as defined above. We use the
MSE for the supervised pixel-wise metric for simplicity but this can be
easily replaced with a perceptual loss if desired.


.. GENERATED FROM PYTHON SOURCE LINES 165-173

.. code-block:: Python


    loss_g = [
        dinv.loss.SupLoss(metric=torch.nn.MSELoss()),
        adversarial.SupAdversarialGeneratorLoss(device=device),
    ]
    loss_d = adversarial.SupAdversarialDiscriminatorLoss(device=device)









.. GENERATED FROM PYTHON SOURCE LINES 174-180

We are now ready to train the networks using :class:`deepinv.training.AdversarialTrainer`.
We load the pretrained models that were trained in the exact same way after 50 epochs,
and fine-tune the model for 1 epoch for a quick demo.
You can find the pretrained models on HuggingFace https://huggingface.co/deepinv/adversarial-demo.
To train from scratch, simply comment out the model loading code and increase the number of epochs.


.. GENERATED FROM PYTHON SOURCE LINES 180-209

.. code-block:: Python


    ckpt = torch.hub.load_state_dict_from_url(
        dinv.models.utils.get_weights_url("adversarial-demo", "deblurgan_model.pth"),
        map_location=lambda s, _: s,
    )

    G.load_state_dict(ckpt["state_dict"])
    D.load_state_dict(ckpt["state_dict_D"])
    optimizer.load_state_dict(ckpt["optimizer"])

    trainer = dinv.training.AdversarialTrainer(
        model=G,
        D=D,
        physics=physics,
        train_dataloader=train_dataloader,
        eval_dataloader=test_dataloader,
        epochs=1,
        losses=loss_g,
        losses_d=loss_d,
        optimizer=optimizer,
        scheduler=scheduler,
        verbose=True,
        show_progress_bar=False,
        save_path=None,
        device=device,
    )

    G = trainer.train()





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

 .. code-block:: none

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

      0%|          | 0.00/12.7M [00:00<?, ?B/s]
     66%|██████▌   | 8.38M/12.7M [00:00<00:00, 87.7MB/s]
    100%|██████████| 12.7M/12.7M [00:00<00:00, 86.1MB/s]
    The model has 444867 trainable parameters
    Train epoch 0: SupLoss=0.004, SupAdversarialGeneratorLoss=0.003, TotalLoss=0.007, PSNR=25.235
    Eval epoch 0: PSNR=25.43




.. GENERATED FROM PYTHON SOURCE LINES 210-212

Test the trained model and plot the results. We compare to the pseudo-inverse as a baseline.


.. GENERATED FROM PYTHON SOURCE LINES 212-217

.. code-block:: Python


    trainer.plot_images = True
    trainer.test(test_dataloader)





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


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

 .. code-block:: none

    Eval epoch 0: PSNR=25.43, PSNR no learning=23.111
    Test results:
    PSNR no learning: 23.111 +- 2.733
    PSNR: 25.430 +- 3.299

    {'PSNR no learning': np.float64(23.110940551757814), 'PSNR no learning_std': np.float64(2.7325177722415908), 'PSNR': np.float64(25.430047607421876), 'PSNR_std': np.float64(3.298856479885389)}



.. GENERATED FROM PYTHON SOURCE LINES 218-236

UAIR training
~~~~~~~~~~~~~

Unsupervised Adversarial Image Reconstruction (UAIR) (Pajot et al.,
`Unsupervised Adversarial Image Reconstruction <https://openreview.net/forum?id=BJg4Z3RqF7>`_)
is a method for solving inverse problems using generative models. In this
example, we use a simple U-Net as the generator and discriminator, and
train using the adversarial loss. The forward pass of the generator is defined as:

**UAIR** forward pass:

.. math:: \hat x = G(y),

**UAIR** loss:

.. math:: \mathcal{L}=\mathcal{L}_\text{adv}(\hat y, y;D)+\lVert \forw{\inverse{\hat y}}- \hat y\rVert^2_2,\quad\hat y=\forw{\hat x}.

We next load the models and construct losses as defined above.

.. GENERATED FROM PYTHON SOURCE LINES 236-245

.. code-block:: Python


    G, D, optimizer, scheduler = get_models(
        lr_g=1e-4, lr_d=4e-4
    )  # learning rates from original paper

    loss_g = adversarial.UAIRGeneratorLoss(device=device)
    loss_d = adversarial.UnsupAdversarialDiscriminatorLoss(device=device)









.. GENERATED FROM PYTHON SOURCE LINES 246-250

We are now ready to train the networks using :class:`deepinv.training.AdversarialTrainer`.
Like above, we load a pretrained model trained in the exact same way for 50 epochs,
and fine-tune here for a quick demo with 1 epoch.


.. GENERATED FROM PYTHON SOURCE LINES 250-278

.. code-block:: Python


    ckpt = torch.hub.load_state_dict_from_url(
        dinv.models.utils.get_weights_url("adversarial-demo", "uair_model.pth"),
        map_location=lambda s, _: s,
    )

    G.load_state_dict(ckpt["state_dict"])
    D.load_state_dict(ckpt["state_dict_D"])
    optimizer.load_state_dict(ckpt["optimizer"])

    trainer = dinv.training.AdversarialTrainer(
        model=G,
        D=D,
        physics=physics,
        train_dataloader=train_dataloader,
        eval_dataloader=test_dataloader,
        epochs=1,
        losses=loss_g,
        losses_d=loss_d,
        optimizer=optimizer,
        scheduler=scheduler,
        verbose=True,
        show_progress_bar=False,
        save_path=None,
        device=device,
    )
    G = trainer.train()





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

 .. code-block:: none

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

      0%|          | 0.00/12.7M [00:00<?, ?B/s]
     18%|█▊        | 2.25M/12.7M [00:00<00:00, 20.9MB/s]
     47%|████▋     | 6.00M/12.7M [00:00<00:00, 29.2MB/s]
     88%|████████▊ | 11.1M/12.7M [00:00<00:00, 37.6MB/s]
    100%|██████████| 12.7M/12.7M [00:00<00:00, 36.8MB/s]
    The model has 444867 trainable parameters
    Train epoch 0: TotalLoss=0.144, PSNR=24.253
    Eval epoch 0: PSNR=23.978




.. GENERATED FROM PYTHON SOURCE LINES 279-281

Test the trained model and plot the results:


.. GENERATED FROM PYTHON SOURCE LINES 281-286

.. code-block:: Python


    trainer.plot_images = True
    trainer.test(test_dataloader)





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


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

 .. code-block:: none

    Eval epoch 0: PSNR=23.978, PSNR no learning=23.111
    Test results:
    PSNR no learning: 23.111 +- 2.733
    PSNR: 23.978 +- 2.705

    {'PSNR no learning': np.float64(23.110940551757814), 'PSNR no learning_std': np.float64(2.7325177722415908), 'PSNR': np.float64(23.978289794921874), 'PSNR_std': np.float64(2.7051143970949734)}



.. GENERATED FROM PYTHON SOURCE LINES 287-318

CSGM / AmbientGAN training
~~~~~~~~~~~~~~~~~~~~~~~~~~

Compressed Sensing using Generative Models (CSGM) and AmbientGAN are two methods for solving inverse problems
using generative models. CSGM (Bora et al., `Compressed Sensing using Generative Models
<https://arxiv.org/abs/1703.03208>`_) uses a generative model to solve the inverse problem by optimising the latent
space of the generator. AmbientGAN (Bora et al., `AmbientGAN: Generative models from lossy measurements
<https://openreview.net/forum?id=Hy7fDog0b>`_) uses a generative model to solve the inverse problem by optimising the
measurements themselves. Both methods are trained using an adversarial loss; the main difference is that CSGM requires
a ground truth dataset (supervised loss), while AmbientGAN does not (unsupervised loss).

In this example, we use a DCGAN as the
generator and discriminator, and train using the adversarial loss. The forward pass of the generator is given by:

**CSGM** forward pass at train time:

.. math:: \hat x = \inverse{z},\quad z\sim \mathcal{N}(\mathbf{0},\mathbf{I}_k)

**CSGM**/**AmbientGAN** forward pass at eval time:

.. math:: \hat x = \inverse{\hat z}\quad\text{s.t.}\quad\hat z=\operatorname*{argmin}_z \lVert \forw{\inverse{z}}-y\rVert _2^2

**CSGM** loss:

.. math:: \mathcal{L}=\mathcal{L}_\text{adv}(\hat x, x;D)

**AmbientGAN** loss (where :math:`\forw{\cdot}` is the physics):

.. math:: \mathcal{L}=\mathcal{L}_\text{adv}(\forw{\hat x}, y;D)

We next load the models and construct losses as defined above.

.. GENERATED FROM PYTHON SOURCE LINES 318-336

.. code-block:: Python


    G = dinv.models.CSGMGenerator(
        dinv.models.DCGANGenerator(output_size=128, nz=100, ngf=32), inf_tol=1e-2
    ).to(device)
    D = dinv.models.DCGANDiscriminator(ndf=32).to(device)
    _, _, optimizer, scheduler = get_models(
        model=G, D=D, lr_g=2e-4, lr_d=2e-4
    )  # learning rates from original paper

    # For AmbientGAN:
    loss_g = adversarial.UnsupAdversarialGeneratorLoss(device=device)
    loss_d = adversarial.UnsupAdversarialDiscriminatorLoss(device=device)

    # For CSGM:
    loss_g = adversarial.SupAdversarialGeneratorLoss(device=device)
    loss_d = adversarial.SupAdversarialDiscriminatorLoss(device=device)









.. GENERATED FROM PYTHON SOURCE LINES 337-344

As before, we can now train our models. Since inference is very
slow for CSGM/AmbientGAN as it requires an optimisation, we only do one
evaluation at the end. Note the train PSNR is meaningless as this
generative model is trained on random latents.
Like above, we load a pretrained model trained in the exact same way for 50 epochs,
and fine-tune here for a quick demo with 1 epoch.


.. GENERATED FROM PYTHON SOURCE LINES 344-372

.. code-block:: Python


    ckpt = torch.hub.load_state_dict_from_url(
        dinv.models.utils.get_weights_url("adversarial-demo", "csgm_model.pth"),
        map_location=lambda s, _: s,
    )

    G.load_state_dict(ckpt["state_dict"])
    D.load_state_dict(ckpt["state_dict_D"])
    optimizer.load_state_dict(ckpt["optimizer"])

    trainer = dinv.training.AdversarialTrainer(
        model=G,
        D=D,
        physics=physics,
        train_dataloader=train_dataloader,
        epochs=1,
        losses=loss_g,
        losses_d=loss_d,
        optimizer=optimizer,
        scheduler=scheduler,
        verbose=True,
        show_progress_bar=False,
        save_path=None,
        device=device,
    )
    G = trainer.train()






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

 .. code-block:: none

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

      0%|          | 0.00/49.3M [00:00<?, ?B/s]
     31%|███       | 15.4M/49.3M [00:00<00:00, 161MB/s]
     69%|██████▉   | 34.1M/49.3M [00:00<00:00, 177MB/s]
    100%|██████████| 49.3M/49.3M [00:00<00:00, 175MB/s]
    The model has 3608000 trainable parameters
    Train epoch 0: TotalLoss=0.007, PSNR=9.136




.. GENERATED FROM PYTHON SOURCE LINES 373-379

Eventually, we run evaluation of the generative model by running test-time optimisation
using test measurements. Note that we do not get great results as CSGM /
AmbientGAN relies on large datasets of diverse samples, and we run the
optimisation to a relatively high tolerance for speed. Improve the results by
running the optimisation for longer.


.. GENERATED FROM PYTHON SOURCE LINES 379-381

.. code-block:: Python


    trainer.test(test_dataloader)




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

 .. code-block:: none

    Eval epoch 0: PSNR=10.25, PSNR no learning=23.111
    Test results:
    PSNR no learning: 23.111 +- 2.733
    PSNR: 10.250 +- 1.120

    {'PSNR no learning': np.float64(23.110940551757814), 'PSNR no learning_std': np.float64(2.7325177722415908), 'PSNR': np.float64(10.24966812133789), 'PSNR_std': np.float64(1.1202718281006854)}




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

   **Total running time of the script:** (1 minutes 25.284 seconds)


.. _sphx_glr_download_auto_examples_adversarial-learning_demo_gan_imaging.py:

.. only:: html

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

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

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

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

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

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

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


.. only:: html

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

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