๐ŸŽ“ A1Learner
๐Ÿ”

PyTorch GAN Failure Modes

GAN failure modes are the two classic ways this two-model game breaks, called mode collapse and a runaway discriminator. This lesson causes both on purpose at small scale and reads each one from its traces.

A certain judge leaves no signal

A runaway discriminator is a judge that becomes too strong too fast. When a judge is already certain, its loss sits near zero. A near-zero loss leaves the generator almost no gradient to learn from.

loss_fn = nn.BCEWithLogitsLoss()
scores = torch.tensor([10.0, -10.0])
print(round(loss_fn(scores, torch.tensor([1.0, 0.0])).item(), 6))

One noise vector, one fake

Mode collapse means the generator repeats one output instead of spreading over the target. A generator fed one fixed noise vector answers with the same output every time, as these three rows show.

torch.manual_seed(11)
gen = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 1))
z_fixed = torch.tensor([[0.5, -1.0, 0.25, 2.0]])
with torch.no_grad():
    out = gen(z_fixed.repeat(3, 1))
print([round(v[0], 2) for v in out.tolist()])

Sampling the collapsed generator

The closing program samples its collapsed generator eight times. The count of different values does the judging.

with torch.no_grad():
    samples = generator(z_fixed.repeat(8, 1))
print(len(set(round(v[0], 2) for v in samples.tolist())))

Example

import torch
from torch import nn

target_mean, target_std = 3.0, 0.5
loss_fn = nn.BCEWithLogitsLoss()


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)


# Part 1: the discriminator trains alone against a frozen generator.
torch.manual_seed(11)
generator = Generator()
discriminator = Discriminator()
opt_d = torch.optim.Adam(discriminator.parameters(), lr=0.02)
first_loss = None
for step in range(400):
    real = torch.randn(64, 1) * target_std + target_mean
    with torch.no_grad():
        fake = generator(torch.randn(64, 4))
    d_loss = (loss_fn(discriminator(real), torch.ones(64, 1))
              + loss_fn(discriminator(fake), torch.zeros(64, 1)))
    if first_loss is None:
        first_loss = d_loss.item()
    opt_d.zero_grad()
    d_loss.backward()
    opt_d.step()

print("Runaway discriminator:")
print("Discriminator loss at step 1:", round(first_loss, 4))
print("Discriminator loss at step 400:", round(d_loss.item(), 4))

# Part 2: the generator only ever sees one fixed noise vector.
torch.manual_seed(11)
generator = Generator()
discriminator = Discriminator()
opt_g = torch.optim.Adam(generator.parameters(), lr=0.005)
opt_d = torch.optim.Adam(discriminator.parameters(), lr=0.005)
z_fixed = torch.randn(1, 4)
for step in range(400):
    real = torch.randn(64, 1) * target_std + target_mean
    fake = generator(z_fixed.repeat(64, 1)).detach()
    d_loss = (loss_fn(discriminator(real), torch.ones(64, 1))
              + loss_fn(discriminator(fake), torch.zeros(64, 1)))
    opt_d.zero_grad()
    d_loss.backward()
    opt_d.step()
    g_loss = loss_fn(discriminator(generator(z_fixed.repeat(64, 1))),
                     torch.ones(64, 1))
    opt_g.zero_grad()
    g_loss.backward()
    opt_g.step()

print("Collapsed generator samples:")
with torch.no_grad():
    samples = generator(z_fixed.repeat(8, 1))
for value in samples.tolist():
    print(round(value[0], 2))
print("Different sample values:", len(set(round(v[0], 2) for v in samples.tolist())))

Key points

  • Mode collapse: the generator repeats one output and ignores the rest of the target
  • Runaway discriminator: the judge's loss falls to near zero early in training
  • A collapsed generator can be read from its samples, which stop varying
  • Training the judge alone against a frozen generator causes the runaway here

In this example, the first part trains a discriminator alone for 400 steps against a frozen generator. Its loss falls from 1.5392 at step 1 to 0.0 at step 400. A judge that certain leaves the generator no useful gradient, which is the runaway failure. The second part starts a fresh pair of models. The generator there only ever sees one fixed noise vector, repeated across every batch. After 400 rounds, eight samples from it all print as 1.07. The count of different sample values is 1. The repeated point also sits away from the target mean of 3.0. The collapse lost both the spread and the place.

Note

The collapse here is caused by holding the noise fixed, which guarantees the repeated output

Try it Yourself

Edit the code, then press Run.

Running is simulated ยท Python Playground

Exercise

A generator's last eight samples are all 1.07. What failure does this show?