class_docstr: ## Configurations
class_docstr_tok: ["#","#","Configurations"]
docstr:         ## Training
  
          We aim to solve:
          $$G^{*}, F^{*} = rg \min_{G,F} \max_{D_X, D_Y} \mathcal{L}(G, F, D_X, D_Y)$$
  
          where,
          $G$ translates images from $X 
  ightarrow Y$,
          $F$ translates images from $Y 
  ightarrow X$,
          $D_X$ tests if images are from $X$ space,
          $D_Y$ tests if images are from $Y$ space, and
  
          egin{align}
          \mathcal{L}(G, F, D_X, D_Y)
              &= \mathcal{L}_{GAN}(G, D_Y, X, Y) \
              &+ \mathcal{L}_{GAN}(F, D_X, Y, X) \
              &+ \lambda_1 \mathcal{L}_{cyc}(G, F) \
              &+ \lambda_2 \mathcal{L}_{identity}(G, F) \
          \
          \mathcal{L}_{GAN}(G, F, D_Y, X, Y)
              &= \mathbb{E}_{y \sim p_{data}(y)} \Big[log D_Y(y)\Big] \
              &+ \mathbb{E}_{x \sim p_{data}(x)} igg[log\Big(1 - D_Y(G(x))\Big)igg] \
              &+ \mathbb{E}_{x \sim p_{data}(x)} \Big[log D_X(x)\Big] \
              &+ \mathbb{E}_{y \sim p_{data}(y)} igg[log\Big(1 - D_X(F(y))\Big)igg] \
          \
          \mathcal{L}_{cyc}(G, F)
              &= \mathbb{E}_{x \sim p_{data}(x)} \Big[\lVert F(G(x)) - x \lVert_1\Big] \
              &+ \mathbb{E}_{y \sim p_{data}(y)} \Big[\lVert G(F(y)) - y 
  Vert_1\Big] \
          \
          \mathcal{L}_{identity}(G, F)
              &= \mathbb{E}_{x \sim p_{data}(x)} \Big[\lVert F(x) - x \lVert_1\Big] \
              &+ \mathbb{E}_{y \sim p_{data}(y)} \Big[\lVert G(y) - y 
  Vert_1\Big] \
          \end{align}
  
          $\mathcal{L}_{GAN}$ is the generative adversarial loss from the original
          GAN paper.
  
          $\mathcal{L}_{cyc}$ is the cyclic loss, where we try to get $F(G(x))$ to be similar to $x$,
          and $G(F(y))$ to be similar to $y$.
          Basically if the two generators (transformations) are applied in series it should give back the
          original image.
          This is the main contribution of this paper.
          It trains the generators to generate an image of the other distribution that is similar to
          the original image.
          Without this loss $G(x)$ could generate anything that's from the distribution of $Y$.
          Now it needs to generate something from the distribution of $Y$ but still has properties of $x$,
          so that $F(G(x)$ can re-generate something like $x$.
  
          $\mathcal{L}_{cyc}$ is the identity loss.
          This was used to encourage the mapping to preserve color composition between
          the input and the output.
  
          To solve $$G^*, F^*$$,
          discriminators $D_X$ and $D_Y$ should **ascend** on the gradient,
  
          egin{align}
          
  abla_{  heta_{D_X, D_Y}} rac{1}{m} \sum_{i=1}^m
          &\Bigg[
          \log D_Y\Big(y^{(i)}\Big) \
          &+ \log \Big(1 - D_Y\Big(G\Big(x^{(i)}\Big)\Big)\Big) \
          &+ \log D_X\Big(x^{(i)}\Big) \
          & +\log\Big(1 - D_X\Big(F\Big(y^{(i)}\Big)\Big)\Big)
          \Bigg]
          \end{align}
  
          That is descend on *negative* log-likelihood loss.
  
          In order to stabilize the training the negative log- likelihood objective
          was replaced by a least-squared loss -
          the least-squared error of discriminator, labelling real images with 1,
          and generated images with 0.
          So we want to descend on the gradient,
  
          egin{align}
          
  abla_{  heta_{D_X, D_Y}} rac{1}{m} \sum_{i=1}^m
          &\Bigg[
              igg(D_Y\Big(y^{(i)}\Big) - 1igg)^2 \
              &+ D_Y\Big(G\Big(x^{(i)}\Big)\Big)^2 \
              &+ igg(D_X\Big(x^{(i)}\Big) - 1igg)^2 \
              &+ D_X\Big(F\Big(y^{(i)}\Big)\Big)^2
          \Bigg]
          \end{align}
  
          We use least-squares for generators also.
          The generators should *descend* on the gradient,
  
          egin{align}
          
  abla_{  heta_{F, G}} rac{1}{m} \sum_{i=1}^m
          &\Bigg[
              igg(D_Y\Big(G\Big(x^{(i)}\Big)\Big) - 1igg)^2 \
              &+ igg(D_X\Big(F\Big(y^{(i)}\Big)\Big) - 1igg)^2 \
              &+ \mathcal{L}_{cyc}(G, F)
              + \mathcal{L}_{identity}(G, F)
          \Bigg]
          \end{align}
  
          We use `generator_xy` for $G$ and `generator_yx` for $F$.
          We use `discriminator_x` for $D_X$ and `discriminator_y` for $D_Y$.
          
code: def run(self):
      """
          ## Training
  
          We aim to solve:
          $$G^{*}, F^{*} = \arg \min_{G,F} \max_{D_X, D_Y} \mathcal{L}(G, F, D_X, D_Y)$$
  
          where,
          $G$ translates images from $X \rightarrow Y$,
          $F$ translates images from $Y \rightarrow X$,
          $D_X$ tests if images are from $X$ space,
          $D_Y$ tests if images are from $Y$ space, and
  
          \begin{align}
          \mathcal{L}(G, F, D_X, D_Y)
              &= \mathcal{L}_{GAN}(G, D_Y, X, Y) \\
              &+ \mathcal{L}_{GAN}(F, D_X, Y, X) \\
              &+ \lambda_1 \mathcal{L}_{cyc}(G, F) \\
              &+ \lambda_2 \mathcal{L}_{identity}(G, F) \\
          \\
          \mathcal{L}_{GAN}(G, F, D_Y, X, Y)
              &= \mathbb{E}_{y \sim p_{data}(y)} \Big[log D_Y(y)\Big] \\
              &+ \mathbb{E}_{x \sim p_{data}(x)} \bigg[log\Big(1 - D_Y(G(x))\Big)\bigg] \\
              &+ \mathbb{E}_{x \sim p_{data}(x)} \Big[log D_X(x)\Big] \\
              &+ \mathbb{E}_{y \sim p_{data}(y)} \bigg[log\Big(1 - D_X(F(y))\Big)\bigg] \\
          \\
          \mathcal{L}_{cyc}(G, F)
              &= \mathbb{E}_{x \sim p_{data}(x)} \Big[\lVert F(G(x)) - x \lVert_1\Big] \\
              &+ \mathbb{E}_{y \sim p_{data}(y)} \Big[\lVert G(F(y)) - y \rVert_1\Big] \\
          \\
          \mathcal{L}_{identity}(G, F)
              &= \mathbb{E}_{x \sim p_{data}(x)} \Big[\lVert F(x) - x \lVert_1\Big] \\
              &+ \mathbb{E}_{y \sim p_{data}(y)} \Big[\lVert G(y) - y \rVert_1\Big] \\
          \end{align}
  
          $\mathcal{L}_{GAN}$ is the generative adversarial loss from the original
          GAN paper.
  
          $\mathcal{L}_{cyc}$ is the cyclic loss, where we try to get $F(G(x))$ to be similar to $x$,
          and $G(F(y))$ to be similar to $y$.
          Basically if the two generators (transformations) are applied in series it should give back the
          original image.
          This is the main contribution of this paper.
          It trains the generators to generate an image of the other distribution that is similar to
          the original image.
          Without this loss $G(x)$ could generate anything that's from the distribution of $Y$.
          Now it needs to generate something from the distribution of $Y$ but still has properties of $x$,
          so that $F(G(x)$ can re-generate something like $x$.
  
          $\mathcal{L}_{cyc}$ is the identity loss.
          This was used to encourage the mapping to preserve color composition between
          the input and the output.
  
          To solve $$G^*, F^*$$,
          discriminators $D_X$ and $D_Y$ should **ascend** on the gradient,
  
          \begin{align}
          \nabla_{\theta_{D_X, D_Y}} \frac{1}{m} \sum_{i=1}^m
          &\Bigg[
          \log D_Y\Big(y^{(i)}\Big) \\
          &+ \log \Big(1 - D_Y\Big(G\Big(x^{(i)}\Big)\Big)\Big) \\
          &+ \log D_X\Big(x^{(i)}\Big) \\
          & +\log\Big(1 - D_X\Big(F\Big(y^{(i)}\Big)\Big)\Big)
          \Bigg]
          \end{align}
  
          That is descend on *negative* log-likelihood loss.
  
          In order to stabilize the training the negative log- likelihood objective
          was replaced by a least-squared loss -
          the least-squared error of discriminator, labelling real images with 1,
          and generated images with 0.
          So we want to descend on the gradient,
  
          \begin{align}
          \nabla_{\theta_{D_X, D_Y}} \frac{1}{m} \sum_{i=1}^m
          &\Bigg[
              \bigg(D_Y\Big(y^{(i)}\Big) - 1\bigg)^2 \\
              &+ D_Y\Big(G\Big(x^{(i)}\Big)\Big)^2 \\
              &+ \bigg(D_X\Big(x^{(i)}\Big) - 1\bigg)^2 \\
              &+ D_X\Big(F\Big(y^{(i)}\Big)\Big)^2
          \Bigg]
          \end{align}
  
          We use least-squares for generators also.
          The generators should *descend* on the gradient,
  
          \begin{align}
          \nabla_{\theta_{F, G}} \frac{1}{m} \sum_{i=1}^m
          &\Bigg[
              \bigg(D_Y\Big(G\Big(x^{(i)}\Big)\Big) - 1\bigg)^2 \\
              &+ \bigg(D_X\Big(F\Big(y^{(i)}\Big)\Big) - 1\bigg)^2 \\
              &+ \mathcal{L}_{cyc}(G, F)
              + \mathcal{L}_{identity}(G, F)
          \Bigg]
          \end{align}
  
          We use `generator_xy` for $G$ and `generator_yx` for $F$.
          We use `discriminator_x` for $D_X$ and `discriminator_y` for $D_Y$.
          """
  
      # Replay buffers to keep generated samples
      gen_x_buffer = ReplayBuffer()
      gen_y_buffer = ReplayBuffer()
  
      # Loop through epochs
      for epoch in monit.loop(self.epochs):
          # Loop through the dataset
          for i, batch in monit.enum("Train", self.dataloader):
              # Move images to the device
              data_x, data_y = batch["x"].to(self.device), batch[
                  "y"
              ].to(self.device)
  
              # true labels equal to $1$
              true_labels = torch.ones(
                  data_x.size(0),
                  *self.discriminator_x.output_shape,
                  device=self.device,
                  requires_grad=False,
              )
              # false labels equal to $0$
              false_labels = torch.zeros(
                  data_x.size(0),
                  *self.discriminator_x.output_shape,
                  device=self.device,
                  requires_grad=False,
              )
  
              # Train the generators.
              # This returns the generated images.
              gen_x, gen_y = self.optimize_generators(
                  data_x, data_y, true_labels
              )
  
              #  Train discriminators
              self.optimize_discriminator(
                  data_x,
                  data_y,
                  gen_x_buffer.push_and_pop(gen_x),
                  gen_y_buffer.push_and_pop(gen_y),
                  true_labels,
                  false_labels,
              )
  
              # Save training statistics and increment the global step counter
              tracker.save()
              tracker.add_global_step(max(len(data_x), len(data_y)))
  
              # Save images at intervals
              batches_done = epoch * len(self.dataloader) + i
              if batches_done % self.sample_interval == 0:
                  # Save models when sampling images
                  experiment.save_checkpoint()
                  # Sample images
                  self.sample_images(batches_done)
  
          # Update learning rates
          self.generator_lr_scheduler.step()
          self.discriminator_lr_scheduler.step()
          # New line
          tracker.new_line()
  
imports: import itertools
  import random
  import zipfile
  from typing import Tuple
  import torch
  import torch.nn
  import torchvision.transforms
  from PIL import Image
  from torch.utils.data import DataLoader, Dataset
  from torchvision.transforms import InterpolationMode
  from torchvision.utils import make_grid
  from labml import lab, tracker, experiment, monit
  from labml.configs import BaseConfigs
  from labml.utils.download import download_file
  from labml.utils.pytorch import get_modules
  from labml_helpers.device import DeviceConfigs
  from labml_helpers.module import Module
label_desc: Use image node_id 3 for calling the Configs obj's underlying member method code with example usage: obj.run() without return types