[0.5.2] - GANs

Colab: exercises | solutions

Please send any problems / bugs on the #errata channel in the Slack group, and ask any questions on the dedicated channels for this chapter of material.

If you want to change to dark mode, you can do this by clicking the three horizontal lines in the top-right, then navigating to Settings → Theme.

Links to all other chapters: (0) Fundamentals, (1) Transformer Interpretability, (2) RL.

Introduction

In this section, we're studying Generative Adversarial Networks (GANs). Together with Variational Autoencoders (VAEs), which are covered in the VAEs section, they're one of the two important classes of generative image models we study in this chapter. Although GANs generally aren't SOTA any more (thanks in part to the rise of diffusion models), studying them teaches you a lot about adversarial training dynamics, and how to debug a model whose loss curves don't tell you whether it's working.

These exercises will also hopefully bring much of this chapter full-circle:

  • We'll be assembling NNet architectures from scratch, like in the ResNets exercises (day 2: CNNs & ResNets)
  • We'll work with different loss functions, and think intuitively about what it means to optimize them (day 3: Optimization & Hyperparameters)
  • We'll have to think carefully about how gradient propagation works between different parts of our model (day 4: Backpropagation)

The generator is built from transposed convolutions. If you haven't done the VAEs section, read its short introduction to transposed convolutions before you start.

As in the VAEs section, the model-building exercises here don't come with rigorous unit tests: all that matters is that your GAN actually trains, and part of the skill we want you to build is reasoning through an architecture from a description or a paper without leaning on a test suite.

For a lecture on this material, which provides some high-level understanding before you dive into the material, watch the video below.

Content & Learning Objectives

1️⃣ GANs

Relative to autoencoders, GANs have a few more moving pieces in their architecture. They're best thought of as two separate networks (the generator and the discriminator) which are learning different goals simultaneously. The goal of the generator is to create images which fool the discriminator, and the goal of the discriminator is to distinguish between real and fake images. The ideal equilibrium point of training is when the generator produces perfect images and the discriminator can't tell the difference between real and fake - however, that's much simpler said than done! GANs are notoriously difficult to train, and we'll have to engage with some of these difficulties during our exercises.

By the end of these exercises, you should have built and trained your own GANs, to generate celebrity pictures. By the time you're done, you'll hopefully have produced output like this (below), and you'll have everything you need to set up a competitor to Midjourney (plus or minus a few other foundational ML papers and 10 million dollars in GPU time).

Learning Objectives
  • Understand the loss function used in GANs, and why it can be expected to result in the generator producing realistic outputs.
  • Implement the DCGAN architecture from the paper, with relatively minimal guidance.
  • Learn how to identify and fix bugs in your GAN architecture, to improve convergence properties.

Setup code

import os
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal

import einops
import torch as t
import torchinfo
import wandb
from einops.layers.torch import Rearrange
from jaxtyping import Float
from torch import Tensor, nn
from torch.utils.data import DataLoader, Dataset, Subset

import torch.nn.functional as F
from tqdm.auto import tqdm


# Make sure exercises are in the path
chapter = "chapter0_fundamentals"
section = "part52_gans"
root_dir = next(p for p in Path.cwd().parents if (p / chapter).exists())
exercises_dir = root_dir / chapter / "exercises"
section_dir = exercises_dir / section
if str(exercises_dir) not in sys.path:
    sys.path.append(str(exercises_dir))

MAIN = __name__ == "__main__"

# Modules from previous exercises (you can swap in your own implementations if you've completed them).
from part2_cnns.solutions import BatchNorm2d, Conv2d, Linear, ReLU, Sequential
from part52_gans import tests
from fundamentals_utils import TANH_RANGE_TRANSFORM, LiveImage, display_data, get_dataset
from plotly_utils import imshow

device = t.device("mps" if t.backends.mps.is_available() else "cuda" if t.cuda.is_available() else "cpu")
NUM_WORKERS = 2 if "google.colab" in sys.modules else min(8, os.cpu_count())
t.set_num_threads(min(4, t.get_num_threads()))