🎓 A1Learner
🔍

PyTorch Training a Tiny GAN

Training a tiny GAN alternates discriminator and generator steps until the generator's spread approaches one target Gaussian. A Gaussian is a bell-shaped spread of values described by a mean and a standard deviation.

Reading a generator's spread

The mean and standard deviation of many samples describe the generator's spread. These six fixed values average 3.05 with a spread of 0.4.

samples = torch.tensor([2.8, 3.4, 2.5, 3.1, 2.9, 3.6])
print(round(samples.mean().item(), 2))
print(round(samples.std().item(), 2))

One round of the two steps

Each round trains the discriminator on fresh real and fake values first. Then one generator step tries to fool the updated discriminator.

real = torch.randn(batch, 1) * target_std + target_mean
fake = generator(torch.randn(batch, 4)).detach()
d_loss = (loss_fn(discriminator(real), torch.ones(batch, 1))
          + loss_fn(discriminator(fake), torch.zeros(batch, 1)))
opt_d.zero_grad()
d_loss.backward()
opt_d.step()

Measuring the end point

Sampling uses its own seed after training ends. The printed mean and spread come from 2,000 fresh fakes.

torch.manual_seed(99)
with torch.no_grad():
    samples = generator(torch.randn(2000, 4))

Example

import torch
from torch import nn

target_mean, target_std = 3.0, 0.5


class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(4, 32), nn.ReLU(),
            nn.Linear(32, 32), nn.ReLU(),
            nn.Linear(32, 1))

    def forward(self, z):
        return self.net(z)


class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, 32), nn.ReLU(),
            nn.Linear(32, 32), nn.ReLU(),
            nn.Linear(32, 1))

    def forward(self, x):
        return self.net(x)


torch.manual_seed(11)
generator = Generator()
discriminator = Discriminator()
opt_g = torch.optim.Adam(generator.parameters(), lr=0.001, betas=(0.5, 0.999))
opt_d = torch.optim.Adam(discriminator.parameters(), lr=0.001, betas=(0.5, 0.999))
loss_fn = nn.BCEWithLogitsLoss()

steps, batch = 1000, 128
for step in range(steps):
    real = torch.randn(batch, 1) * target_std + target_mean
    fake = generator(torch.randn(batch, 4)).detach()
    d_loss = (loss_fn(discriminator(real), torch.ones(batch, 1))
              + loss_fn(discriminator(fake), torch.zeros(batch, 1)))
    opt_d.zero_grad()
    d_loss.backward()
    opt_d.step()

    g_loss = loss_fn(discriminator(generator(torch.randn(batch, 4))),
                     torch.ones(batch, 1))
    opt_g.zero_grad()
    g_loss.backward()
    opt_g.step()

torch.manual_seed(99)
with torch.no_grad():
    samples = generator(torch.randn(2000, 4))

print("Target mean:", target_mean, " Target spread (std):", target_std)
print("Generated mean:", round(samples.mean().item(), 2))
print("Generated spread (std):", round(samples.std().item(), 2))
print("Final discriminator loss:", round(d_loss.item(), 2))
print("Final generator loss:", round(g_loss.item(), 2))

Key points

  • The target here is a mean of 3.0 with a standard deviation of one half
  • The generator's mean lands on the target mean after 1,000 rounds
  • Its spread reaches only part of the target width in this run
  • Both losses settle near the values of a judge guessing at chance

In this example, a seed builds both models. Adam trains them with a low rate and a short gradient memory. The betas argument of Adam sets that memory, and the pair (0.5, 0.999) steadies this two-model game. Each of the 1,000 rounds uses a batch of 128 fresh values. The discriminator steps first, and the generator steps second. The final discriminator loss prints as 1.39, near twice the loss of a guessing judge. The final generator loss prints as 0.7, near the loss of a judge it fully fools half the time. A seeded sample of 2,000 fakes then gives the end point. The generated mean is 2.99 against a target of 3.0. The generated spread is 0.33 against a target of 0.5. So the fakes sit in a narrower band than the real values.

Note

GAN endpoints vary with the seed, so this lesson states the values its own seed produces

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

After training, a generator's samples average 2.99 against a target mean of 3.0, with spread 0.33 against a target of 0.5. Which part matched best?