GAN Training Loop: Separate Updates

Hello! Welcome back to our journey into generative models.

In our last lesson, we designed the architectures for our two key players: the Generator (the art forger) and the Discriminator (the art detective). With these models defined in PyTorch, we now have the components ready.

Today, we'll bring them into the training arena. Our goal is to implement the GAN training loop with separate updates for the generator and discriminator. This isn't a typical training process with a single loss function. Instead, we'll orchestrate the two-player minimax game where each network learns from the other in a carefully choreographed dance.

1. Setting the Stage for Training

Before we start the loop, we need to set up our training environment. This involves defining the loss function and, crucially, separate optimizers for our two networks.

Loss Function and Labels

For a GAN's binary classification task (real vs. fake), the Binary Cross Entropy (BCE) Loss is the perfect tool. In PyTorch, this is nn.BCELoss. The loss for a single prediction is given by:

Here, $x$ is the Discriminator's output (a probability) and $y$ is the ground-truth label. By setting the label $y$ to 1 for real images and 0 for fake images, we can use this single function to calculate the loss for both scenarios.

Two Separate Optimizers

Since the Generator and Discriminator are two distinct networks with different objectives, they must be trained by two separate optimizers. One optimizer will update the Discriminator's parameters, and the other will update only the Generator's parameters. This is the core of implementing separate updates. We'll use the Adam optimizer for both, as recommended in the original DCGAN paper.

2. The Adversarial Training Loop

The training loop proceeds for a number of epochs, and within each epoch, it iterates through batches of our real data. For each batch, we perform a two-part update: first, we train the Discriminator, and second, we train the Generator.

Now, let's see how this is implemented step-by-step.

3. Part 1: Train the Discriminator

The Discriminator's goal is to get better at distinguishing real images from fake ones. Its training involves two mini-steps: learning from a batch of real images and learning from a batch of fake images.

Let's summarize the steps to train the Discriminator for one batch:

1. Zero the gradients: Clear out any old gradients from the previous iteration.
   optimizerD.zero_grad()

2. Train on real images:

   • Take a batch of real_images from your dataloader.
   • Pass them through the Discriminator: output_real = netD(real_images).
   • Calculate the loss by comparing the output to "real" labels (a tensor of all 1s): lossD_real = criterion(output_real, real_labels).
   • Compute the gradients for this loss: lossD_real.backward().

3. Train on fake images:

   • Generate a batch of fake_images using the Generator.
   • Pass these fakes through the Discriminator. Crucially, you must detach() the fake images from the Generator's computational graph: output_fake = netD(fake_images.detach()).
   • Why .detach()? We are only training the Discriminator right now. By detaching, we tell PyTorch to treat the fake images as fixed inputs (like an image loaded from a file). This prevents gradients from this step from flowing backward into the Generator, which would wrongly update the Generator based on the Discriminator's loss.
   • Calculate the loss by comparing the output to "fake" labels (a tensor of all 0s): lossD_fake = criterion(output_fake, fake_labels).
   • Compute and accumulate the gradients for this loss: lossD_fake.backward().

4. Update the Discriminator:

   • The total loss is the sum of the real and fake losses: lossD = lossD_real + lossD_fake.
   • Now, step the Discriminator's optimizer to update its weights using the accumulated gradients: optimizerD.step().

The detective is now a little bit smarter.

4. Part 2: Train the Generator

The Generator's goal is to produce fakes that are so good they fool the (now smarter) Discriminator.

The original GAN paper proposed that the Generator should try to minimize the probability that its fakes are caught, which corresponds to minimizing $\log (1-D(G(z)))$. However, this objective can lead to weak gradients early in training. A more effective and widely-used approach is to train the Generator to maximize the probability that its fakes are classified as real, i.e., maximize $\log (D(G(z)))$.

Here are the steps to train the Generator:

1. Zero the gradients: Clear the Generator's old gradients.
   optimizerG.zero_grad()

2. Classify the fake images:

   • Take the fake_images we generated in the previous part and pass them through the updated Discriminator: output = netD(fake_images).
   • No .detach() this time! We need the gradients to be able to flow from this output, back through the Discriminator's layers (without changing them), and all the way back into the Generator's layers so it can learn.

3. Calculate Generator's loss:

   • This is the clever part. We calculate the Generator's loss by comparing the Discriminator's output on the fake images to real labels (a tensor of all 1s): lossG = criterion(output, real_labels).
   • This effectively "lies" to the loss function. We are telling the Generator its goal is to produce images that the Discriminator classifies as 1 (real). The resulting loss lossG measures how well the Generator failed at this goal.

4. Update the Generator:

   • Compute the gradients that will teach the Generator how to better fool the Discriminator: lossG.backward().
   • Step the Generator's optimizer to update its weights: optimizerG.step().

The forger is now a little more sophisticated.

5. The Complete Training Loop in Code

By placing these two parts inside a for loop that iterates over our dataset, we have our complete GAN training algorithm.

For a more advanced DCGAN example, the PyTorch tutorial provides the full code. The logic is identical, just with more complex models.

Conclusion

In this lesson, you've learned how to construct the heart of a GAN: the adversarial training loop. This is a significant step, as the interplay between the two networks is what makes GANs so powerful and unique.

Key Takeaways:

• GAN training requires two separate optimizers, one for the Generator and one for the Discriminator.
• Each training iteration has two parts: updating the Discriminator and updating the Generator.
• When training the Discriminator, you must .detach() the fake images from the Generator's graph to prevent its gradients from flowing backward.
• When training the Generator, you calculate its loss against real labels (1s), effectively teaching it to produce images that the Discriminator classifies as real.

Preview of the Next Lesson

We have now defined the networks and implemented the training logic. All the pieces are in place. In the next lesson, we will finally train a GAN, generate synthetic images, and compare GANs vs. VAEs in terms of their output quality and training characteristics. We'll run the code we've prepared and analyze the results of our adversarial game.