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

#### 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)..

<a id="sphx-glr-auto-examples-adversarial-learning-demo-gan-imaging-py"></a>

# 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 [`deepinv.training.AdversarialTrainer`](https://deepinv.org/api/stubs/deepinv.training.AdversarialTrainer.html.md#deepinv.training.AdversarialTrainer). These
examples can also be easily extended to train more complicated GANs such
as CycleGAN.

This example is based on the papers DeblurGAN <sup>[1](#footcite-kupyn2018deblurgan)</sup>,
Compressed Sensing using Generative Models (CSGM) <sup>[2](#footcite-bora2017compressed)</sup>,
AmbientGAN <sup>[3](#footcite-bora2018ambientgan)</sup>, and Unsupervised Adversarial Image Reconstruction (UAIR) <sup>[4](#footcite-pajot2019unsupervised)</sup>.

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

$$
\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 $D(\cdot)$ is the discriminator model, $x$ is the
reference image, $\hat x$ is the estimated reconstruction,
$q(\cdot)$ is a quality function (e.g $q(x)=x$ for WGAN).
Training alternates between generator $G$ and discriminator
$D$ in a minimax game. When there are no ground truths (i.e.
unsupervised), this may be defined on the measurements $y$
instead.

<!-- GENERATED FROM PYTHON SOURCE LINES 31-51 -->
```Python
from pathlib import Path

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

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

device = dinv.utils.get_device()

BASE_DIR = Path(".")
DATA_DIR = BASE_DIR / "measurements"
ORIGINAL_DATA_DIR = get_cache_home() / "datasets" / "Urban100"
```

```none
Selected GPU 0 with 15239.25 MiB free memory
```

<!-- GENERATED FROM PYTHON SOURCE LINES 53-60 -->

## Generate dataset

In this example we use the Urban100 dataset resized to 128x128. We apply random
motion blur physics using
[`deepinv.physics.generator.MotionBlurGenerator`](https://deepinv.org/api/stubs/deepinv.physics.generator.MotionBlurGenerator.html.md#deepinv.physics.generator.MotionBlurGenerator), and save the data
using [`deepinv.datasets.generate_dataset()`](https://deepinv.org/api/stubs/deepinv.datasets.generate_dataset.html.md#deepinv.datasets.generate_dataset).

<!-- GENERATED FROM PYTHON SOURCE LINES 60-103 -->
```Python
physics = dinv.physics.Blur(padding="circular", device=device)
blur_generator = MotionBlurGenerator((11, 11), device=device)

dataset = dinv.datasets.Urban100HR(
    root=ORIGINAL_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,
        load_physics_generator_params=True,
        use_dict_output=True,
    ),
    shuffle=True,
)
test_dataloader = DataLoader(
    dinv.datasets.HDF5Dataset(
        dataset_path,
        train=False,
        load_physics_generator_params=True,
        use_dict_output=True,
    ),
    shuffle=False,
)
```

```none
/local/jtachell/deepinv/deepinv/deepinv/datasets/base.py:495: DeprecationWarning: The tuple format for dataset outputs is deprecated and will be removed in a future version.It is recommended to set `use_dict_output=True` for better readability and flexibility in returned outputs.The default is currently `False` for backward compatibility, but will be switched to `True` in a future version.
  super().__init__(use_dict_output=use_dict_output)
Dataset has been saved at measurements/dinv_dataset0.h5
```

<!-- GENERATED FROM PYTHON SOURCE LINES 104-114 -->

## 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 <sup>[5](#footcite-isola2017image)</sup>, but
these can be replaced with any architecture e.g transformers, unrolled
etc. Further discriminator models are in [adversarial models](https://deepinv.org/user_guide/reconstruction/adversarial.html.md#adversarial).

<!-- GENERATED FROM PYTHON SOURCE LINES 114-141 -->
```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 142-161 -->

## Conditional GAN training

Conditional GANs <sup>[1](#footcite-kupyn2018deblurgan)</sup> 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:

$$
\hat x = G(y)

$$

**Conditional GAN** loss:

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

$$

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

<!-- GENERATED FROM PYTHON SOURCE LINES 161-165 -->
```Python
G, D, optimizer, scheduler = get_models()
```

<!-- GENERATED FROM PYTHON SOURCE LINES 166-170 -->

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 170-178 -->
```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 179-185 -->

We are now ready to train the networks using [`deepinv.training.AdversarialTrainer`](https://deepinv.org/api/stubs/deepinv.training.AdversarialTrainer.html.md#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](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 185-214 -->
```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()
```

```none
/local/jtachell/deepinv/deepinv/deepinv/training/adversarial.py:165: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
  super().setup_train(**kwargs)
The model has 444867 trainable parameters
/local/jtachell/deepinv/deepinv/deepinv/training/adversarial.py:168: UserWarning: optimizer_step_multi_dataset parameter of Trainer is should be set to `False` when using adversarial trainer. Automatically setting it to `False`.
  warnings.warn(
Train epoch 0: SupLoss=0.005, SupAdversarialGeneratorLoss=0.003, TotalLoss=0.008, PSNR=24.785
Eval epoch 0: PSNR=26.728
Best model saved at epoch 1
```

<!-- GENERATED FROM PYTHON SOURCE LINES 215-217 -->

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

<!-- GENERATED FROM PYTHON SOURCE LINES 217-222 -->
```Python
trainer.plot_images = True
trainer.test(test_dataloader)
```

```none
Eval epoch 0: PSNR=26.728, PSNR no learning=24.984
Test results:
PSNR no learning: 24.984 +- 4.028
PSNR: 26.728 +- 3.310

{'PSNR no learning': 24.98430299758911, 'PSNR no learning_std': 4.02772183470814, 'PSNR': 26.727653121948244, 'PSNR_std': 3.3100562181070288, 'runtime': 0.06240446912124753, 'peak_gpu_memory_usage': 0.8108348846435547}
```

<!-- GENERATED FROM PYTHON SOURCE LINES 223-240 -->

## UAIR training

Unsupervised Adversarial Image Reconstruction (UAIR) <sup>[4](#footcite-pajot2019unsupervised)</sup>
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:

$$
\hat x = G(y),

$$

**UAIR** loss:

$$
\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 240-249 -->
```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 250-254 -->

We are now ready to train the networks using [`deepinv.training.AdversarialTrainer`](https://deepinv.org/api/stubs/deepinv.training.AdversarialTrainer.html.md#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 254-282 -->
```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()
```

```none
/local/jtachell/deepinv/deepinv/deepinv/training/adversarial.py:165: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
  super().setup_train(**kwargs)
The model has 444867 trainable parameters
/local/jtachell/deepinv/deepinv/deepinv/training/adversarial.py:168: UserWarning: optimizer_step_multi_dataset parameter of Trainer is should be set to `False` when using adversarial trainer. Automatically setting it to `False`.
  warnings.warn(
Train epoch 0: TotalLoss=0.138, PSNR=24.231
Eval epoch 0: PSNR=25.383
Best model saved at epoch 1
```

<!-- GENERATED FROM PYTHON SOURCE LINES 283-285 -->

Test the trained model and plot the results:

<!-- GENERATED FROM PYTHON SOURCE LINES 285-290 -->
```Python
trainer.plot_images = True
trainer.test(test_dataloader)
```

```none
Eval epoch 0: PSNR=25.383, PSNR no learning=24.984
Test results:
PSNR no learning: 24.984 +- 4.028
PSNR: 25.383 +- 2.934

{'PSNR no learning': 24.98430299758911, 'PSNR no learning_std': 4.02772183470814, 'PSNR': 25.38293113708496, 'PSNR_std': 2.93398121906251, 'runtime': 0.0716500012204051, 'peak_gpu_memory_usage': 0.8108348846435547}
```

<!-- GENERATED FROM PYTHON SOURCE LINES 291-320 -->

## CSGM / AmbientGAN training

Compressed Sensing using Generative Models (CSGM) <sup>[2](#footcite-bora2017compressed)</sup> and AmbientGAN <sup>[3](#footcite-bora2018ambientgan)</sup> are two methods for solving inverse problems
using generative models. CSGM uses a generative model to solve the inverse problem by optimising the latent
space of the generator. AmbientGAN 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:

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

$$

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

$$
\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:

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

$$

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

$$
\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 320-338 -->
```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 339-346 -->

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 346-374 -->
```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()
```

```none
/local/jtachell/deepinv/deepinv/deepinv/training/adversarial.py:165: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
  super().setup_train(**kwargs)
The model has 3608000 trainable parameters
/local/jtachell/deepinv/deepinv/deepinv/training/adversarial.py:168: UserWarning: optimizer_step_multi_dataset parameter of Trainer is should be set to `False` when using adversarial trainer. Automatically setting it to `False`.
  warnings.warn(
Train epoch 0: TotalLoss=0.007, PSNR=9.123
```

<!-- GENERATED FROM PYTHON SOURCE LINES 375-381 -->

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 381-384 -->
```Python
trainer.test(test_dataloader)
```

```none
Eval epoch 0: PSNR=9.506, PSNR no learning=24.984
Test results:
PSNR no learning: 24.984 +- 4.028
PSNR: 9.506 +- 1.380

{'PSNR no learning': 24.98430299758911, 'PSNR no learning_std': 4.02772183470814, 'PSNR': 9.505514335632324, 'PSNR_std': 1.3803674545242581, 'runtime': 0.5724361767061055, 'peak_gpu_memory_usage': 0.8026571273803711}
```

<!-- GENERATED FROM PYTHON SOURCE LINES 385-388 -->
* **References:**

* <a id='footcite-kupyn2018deblurgan'>**[1]**</a> Orest Kupyn, Volodymyr Budzan, Mykola Mykhailych, Dmytro Mishkin, and Jiří Matas. Deblurgan: blind motion deblurring using conditional adversarial networks. In *Proceedings of the IEEE conference on computer vision and pattern recognition*, 8183–8192. 2018.
* <a id='footcite-bora2017compressed'>**[2]**</a> Ashish Bora, Ajil Jalal, Eric Price, and Alexandros G Dimakis. Compressed sensing using generative models. In *International conference on machine learning*, 537–546. PMLR, 2017.
* <a id='footcite-bora2018ambientgan'>**[3]**</a> Ashish Bora, Eric Price, and Alexandros G Dimakis. Ambientgan: generative models from lossy measurements. In *International conference on learning representations*. 2018.
* <a id='footcite-pajot2019unsupervised'>**[4]**</a> Arthur Pajot, Emmanuel De Bézenac, and Patrick Gallinari. Unsupervised adversarial image reconstruction. In *International conference on learning representations*. 2019.
* <a id='footcite-isola2017image'>**[5]**</a> Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In *Proceedings of the IEEE conference on computer vision and pattern recognition*, 1125–1134. 2017.

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

🏷 Tags: Deblurring

<a id="sphx-glr-download-auto-examples-adversarial-learning-demo-gan-imaging-py"></a>

[`Download Jupyter notebook: demo_gan_imaging.ipynb`](https://deepinv.org/auto_examples/adversarial-learning/_downloads/2f3e20ed263cc3e5c7e770603c2bda28/demo_gan_imaging.ipynb)

[`Download Python source code: demo_gan_imaging.py`](https://deepinv.org/auto_examples/adversarial-learning/_downloads/1b08a6c94b052657775204a9369c2c98/demo_gan_imaging.py)

[`Download zipped: demo_gan_imaging.zip`](https://deepinv.org/auto_examples/adversarial-learning/_downloads/74982483541b96cb559b9926a0156344/demo_gan_imaging.zip)

[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)