---
title: "Generative Deep Learning: Generative adversarial networks"
output: 
  html_notebook: 
    theme: cerulean
    highlight: textmate
editor_options: 
  chunk_output_type: console
---

<!-- knitr::purl(input = "static/slides/code/ch12_GAN05_Generative_adversarial_network.Rmd", output = "static/slides/code/ch12_GAN05_Generative_adversarial_network.R",documentation = 0) -->

```{r}
library(keras)
library(tensorflow)
library(tfdatasets)
library(tfautograph)
```
```{r}
# Load, resize, and format pictures into appropriate arrays
preprocess_image <- tf_function(function(image_path) {
  image_path %>%
    tf$io$read_file() %>%
    tf$io$decode_image() %>%
    # Add batch axis, equivalent to .[tf$newaxis, all_dims()]. axis arg is 0-based.
    tf$expand_dims(axis = 0L) %>%
    # Cast from 'uint8'.
    tf$cast("float32") %>%
    inception_v3_preprocess_input()
})

# Convert a tensor array into a valid image and undo preprocessing 
deprocess_image <- tf_function(function(img) {
  img %>%
    # Drop first dim—the batch axis (must be size 1), the inverse of tf$expand_dims().
    tf$squeeze(axis = 0L) %>%
    # Rescale so values in [-1, 1] are remapped to [0, 255].
    { (. * 127.5) + 127.5 } %>%
    # saturate_case() clips values to the dtype range: [0, 255].
    tf$saturate_cast("uint8")
})


display_image_tensor <- function(x, ..., max = 255,
                                 plot_margins = c(0, 0, 0, 0)) {

  if (!is.null(plot_margins))
    # Default to no margins when plotting images.
    withr::local_par(mar = plot_margins)

  x %>%
    # Convert tensors to R arrays.
    as.array() %>%
    drop() %>%
    # Convert to R native raster format.
    as.raster(max = max) %>%
    plot(..., interpolate = FALSE)
}
```

## 12.5 Introduction to generative adversarial networks

**Generative Adversarial Networks (GANs)** are another approach, besides Variational Autoencoders (VAEs), for **learning latent spaces of images**. Introduced in 2014, GANs aim to generate **highly realistic synthetic images** that are statistically almost indistinguishable from real images in a training dataset. The concept is illustrated through an analogy of a **forger trying to create fake art and an expert trying to detect the fakes**.

### 12.5.1 A schematic GAN implementation

A GAN consists of two main neural networks: a **generator** and a **discriminator** (or adversary). The **generator** takes a random vector from a latent space as input and outputs a synthetic image. The **discriminator** takes an image (either real from the dataset or synthetic from the generator) as input and predicts whether it's real or fake.

The training process is adversarial: the **generator is trained to fool the discriminator** into believing its generated images are real, while the **discriminator is trained to correctly identify real and fake images**. Through this process, the generator evolves to produce increasingly realistic images. Once training is complete, the generator can take any point in its input (latent) space and generate a believable image. However, unlike VAEs, the latent space learned by GANs may **lack the same level of explicit structure and continuity**. Finally, the section notes that training GANs is **notoriously difficult** due to its dynamic nature of seeking an equilibrium between two competing networks, rather than minimizing a fixed loss function.

### 12.5.2 A bag of tricks

Training GANs and tuning their implementations is notoriously difficul. Unlike standard deep learning training with a fixed loss landscape, GAN training involves a dynamic system where two networks (generator and discriminator) are constantly competing, seeking an equilibrium rather than a minimum. This makes getting a GAN to work correctly require **lots of careful tuning of the model architecture and training parameters**.

*   Using **strides instead of pooling for downsampling feature maps in the discriminator**, similar to the VAE encoder.
*   **Sampling points from the latent space using a normal (Gaussian) distribution**, rather than a uniform distribution.
*   Introducing **stochasticity by adding random noise to the labels for the discriminator**, which can help prevent the GAN from getting stuck.
*   **Avoiding sparse gradients**, which can hinder GAN training, by:
    *   Using **strided convolutions for downsampling** instead of max pooling.
    *   Using **`layer_activation_leaky_relu()` instead of ReLU activation**, as LeakyReLU allows small negative activation values, relaxing sparsity constraints.
*   Addressing **checkerboard artifacts** in generated images (caused by unequal pixel space coverage) by ensuring that whenever a strided `layer_conv_2d_transpose()` or `layer_conv_2d()` is used in both the generator and the discriminator, the **kernel size is divisible by the stride size**.

### 12.5.3 Getting our hands on the CelebA dataset

The **CelebFaces Attributes Dataset (CelebA)** is a **large-scale dataset of 200,000 faces of celebrities**. We will use this dataset to train a Deep Convolutional GAN (DCGAN) to generate images of human faces. 


```{r eval=FALSE}
# Install gdown.
reticulate::py_install("gdown", pip = TRUE)
# Download the compressed data using gdown.
system("gdown 1O7m1010EJjLE5QxLZiM9Fpjs7Oj6e684")
# Uncompress the data.
zip::unzip("img_align_celeba.zip", exdir = "celeba_gan")

```

```{r}
dataset <- image_dataset_from_directory(
  "celeba_gan",
  # Only the images will be returned—no labels.
  label_mode = NULL,
  # We will resize the images to 64 × 64 by using a smart combination of cropping 
  # and resizing to preserve aspect ratio. We don’t want face proportions to get distorted!
  image_size = c(64, 64),
  batch_size = 32,
  crop_to_aspect_ratio = TRUE
)

# rescale the images to the [0-1] range
dataset %<>% dataset_map(~ .x / 255)

# Displaying the first image
x <- dataset %>% as_iterator() %>% iter_next()
display_image_tensor(x[1, , , ], max = 1)

```

### 12.5.4 The discriminator

We define the **discriminator network** for our Generative Adversarial Network (GAN), which we are training on the **CelebA dataset** of celebrity faces. The discriminator's role is to **take an input image** (either a real image from the dataset or a fake image generated by the generator) and **classify it as either "real" or "generated"**.

It consists of several **convolutional layers (`layer_conv_2d`)** that learn to extract features from the input image. These convolutional layers use a **kernel size of 4** and a **stride of 2** for downsampling (preferred over pooling). Following each convolutional layer (except the last), a **Leaky ReLU activation function (`layer_activation_leaky_relu`)** with an alpha of 0.2 is applied. The use of **Leaky ReLU** instead of standard ReLU is another trick to **avoid sparse gradients**, which can hinder GAN training. Leaky ReLU allows small negative activation values.

After the convolutional layers, the feature maps are **flattened (`layer_flatten`)** into a 1D vector. A **dropout layer (`layer_dropout`)** with a rate of 0.2 is then added. **Dropout** is a trick used in the discriminator to help **prevent the generator from getting stuck** with generating noisy images. Finally, a **dense layer (`layer_dense`)** with a single output unit and a **sigmoid activation function** is used to produce a probability score between 0 and 1, representing the discriminator's confidence that the input image is real (closer to 1) or fake (closer to 0).

```{r}
discriminator <-
  keras_model_sequential(name = "discriminator",
                         input_shape = c(64, 64, 3)) %>%
  layer_conv_2d(64, kernel_size = 4, strides = 2, padding = "same") %>%
  layer_activation_leaky_relu(alpha = 0.2) %>%
  layer_conv_2d(128, kernel_size = 4, strides = 2, padding = "same") %>%
  layer_activation_leaky_relu(alpha = 0.2) %>%
  layer_conv_2d(128, kernel_size = 4, strides = 2, padding = "same") %>%
  layer_activation_leaky_relu(alpha = 0.2) %>%
  layer_flatten() %>%
  # One dropout layer: an important trick!
  layer_dropout(0.2) %>%
  layer_dense(1, activation = "sigmoid")

discriminator
```

### 12.5.5 The generator

Next, we define the **generator network** for our Generative Adversarial Network (GAN), which is tasked with **creating synthetic images of celebrity faces** from random noise. The generator takes a **low-dimensional vector from the latent space** as input and transforms it into a **64x64 pixel RGB image**.

The generator network takes a 128-dimensional random vector, expands it through a dense layer, reshapes it into a small feature map, and then uses a series of transposed convolutional layers with Leaky ReLU activations to upsample it to a 64x64 RGB image with pixel values between 0 and 1. The architecture incorporates tricks like using a kernel size divisible by the stride in transposed convolutions and Leaky ReLU activations for more stable and better-quality image generation.

```{r}
# The latent space will be made of 128-dimensional vectors.
latent_dim <- 128

generator <-
  keras_model_sequential(name = "generator",
                         input_shape = c(latent_dim)) %>%
  # Produce the same number of coefficients we had at the level of
  # the Flatten layer in the encoder.
  layer_dense(8 * 8 * 128) %>%
  # Revert the layer_flatten() of the encoder.
  layer_reshape(c(8, 8, 128)) %>%
  # Revert the layer_conv_2d() of the encoder.
  layer_conv_2d_transpose(128, kernel_size = 4, strides = 2,
                          padding = "same") %>%
  layer_activation_leaky_relu(alpha = 0.2) %>%
  layer_conv_2d_transpose(256, kernel_size = 4, strides = 2,
                          padding = "same") %>%
  # Use Leaky Relu as our activation
  layer_activation_leaky_relu(alpha = 0.2) %>%
  layer_conv_2d_transpose(512, kernel_size = 4, strides = 2,
                          padding = "same") %>%
  layer_activation_leaky_relu(alpha = 0.2) %>%
  layer_conv_2d(3, kernel_size = 5, padding = "same",
                activation = "sigmoid")

generator
```

### 12.5.6 The adversarial network

Here, we define the **`GAN` model class**, which **combines the previously defined discriminator and generator networks**. This combined model is crucial for training the generator to produce realistic images that can fool the discriminator.

To recapitulate, this is what the training loop looks like schematically. For each
epoch, you do the following:

1. Draw random points in the latent space (random noise).
2. Generate images with generator using this random noise.
3. Mix the generated images with real ones.
4. Train discriminator using these mixed images, with corresponding targets:
either “real” (for the real images) or “fake” (for the generated images).
5. Draw new random points in the latent space.
6. Train generator using these random vectors, with targets that all say “these are
real images.” This updates the weights of the generator to move them toward
getting the discriminator to predict “these are real images” for generated
images: this trains the generator to fool the discriminator.

```{r}
model_gan <- new_model_class(
  classname = "GAN",

  initialize = function(discriminator, generator, latent_dim) {
    super$initialize()
    self$discriminator  <- discriminator
    self$generator      <- generator
    self$latent_dim     <- as.integer(latent_dim)
    # Set up metrics to track the two losses over each training epoch
    self$d_loss_metric  <- metric_mean(name = "d_loss")
    self$g_loss_metric  <- metric_mean(name = "g_loss")
  },

  compile = function(d_optimizer, g_optimizer, loss_fn) {
    super$compile()
    self$d_optimizer <- d_optimizer
    self$g_optimizer <- g_optimizer
    self$loss_fn <- loss_fn
  },

   metrics = mark_active(function() {
      list(self$d_loss_metric,
           self$g_loss_metric)
   }),

  # train_step is called with a batch of real images.
  train_step = function(real_images) {
    batch_size <- tf$shape(real_images)[1]
    # Sample random points in the latent space.
    random_latent_vectors <-
      tf$random$normal(shape = c(batch_size, self$latent_dim))
    # Decode them to fake images.
    generated_images <- self$generator(random_latent_vectors)

    # Combine them with real images.
    combined_images <-
      tf$concat(list(generated_images,
                     real_images),
                axis = 0L)

    # Assemble labels, discriminating real from fake images.
    labels <-
      tf$concat(list(tf$ones(tuple(batch_size, 1L)),
                     tf$zeros(tuple(batch_size, 1L))),
                axis = 0L)

    # Add random noise to the labels—an important trick!
    labels %<>% `+`(
      tf$random$uniform(tf$shape(.), maxval = 0.05))

    with(tf$GradientTape() %as% tape, {
      predictions <- self$discriminator(combined_images)
      d_loss <- self$loss_fn(labels, predictions)
    })

    grads <- tape$gradient(d_loss, self$discriminator$trainable_weights)
    # Train the discriminator.
    self$d_optimizer$apply_gradients(
      zip_lists(grads, self$discriminator$trainable_weights))

    # Sample random points in the latent space.
    random_latent_vectors <-
      tf$random$normal(shape = c(batch_size, self$latent_dim))

    # Assemble labels that say “these are all real images” (it’s a lie!).
    misleading_labels <- tf$zeros(tuple(batch_size, 1L))

    with(tf$GradientTape() %as% tape, {
      predictions <- random_latent_vectors %>%
        self$generator() %>%
        self$discriminator()
      g_loss <- self$loss_fn(misleading_labels, predictions)
    })
    grads <- tape$gradient(g_loss, self$generator$trainable_weights)
    # Train the generator.
    self$g_optimizer$apply_gradients(
      zip_lists(grads, self$generator$trainable_weights))

    self$d_loss_metric$update_state(d_loss)
    self$g_loss_metric$update_state(g_loss)

    list(d_loss = self$d_loss_metric$result(),
         g_loss = self$g_loss_metric$result())
  })
```

A callback that samples generated images during training

```{r}
callback_gan_monitor <- new_callback_class(
  classname = "GANMonitor",

  initialize = function(num_img = 3, latent_dim = 128,
                        dirpath = "gan_generated_images") {
    private$num_img <- as.integer(num_img)
    private$latent_dim <- as.integer(latent_dim)
    private$dirpath <- fs::path(dirpath)
    fs::dir_create(dirpath)
  },

  on_epoch_end = function(epoch, logs = NULL) {
    random_latent_vectors <-
      tf$random$normal(shape = c(private$num_img, private$latent_dim))

    generated_images <- random_latent_vectors %>%
      self$model$generator() %>%
      # Scale and clip to uint8 range of [0, 255], and cast to uint8.
      { tf$saturate_cast(. * 255, "uint8") }

    for (i in seq(private$num_img))
      tf$io$write_file(
        filename = private$dirpath / sprintf("img_%03i_%02i.png", epoch, i),
        contents = tf$io$encode_png(generated_images[i, , , ])
      )
  }
)
```

Compiling and training the GAN

```{r}
epochs <- 100

gan <- model_gan(discriminator = discriminator,
                 generator = generator,
                 latent_dim = latent_dim)

gan %>% compile(
  d_optimizer = optimizer_adam(learning_rate = 0.0001),
  g_optimizer = optimizer_adam(learning_rate = 0.0001),
  loss_fn = loss_binary_crossentropy()
)


gan %>% fit(
  dataset,
  epochs = epochs,
  callbacks = callback_gan_monitor(num_img = 10, latent_dim = latent_dim)
)

```
