Generative Adversarial Networks (GAN)

import math, torch, torch.nn as nn

torch.manual_seed(0)
# Real data: 8 Gaussian blobs on a circle, a classic test for "mode collapse"
centers = torch.tensor([[math.cos(2 * math.pi * k / 8), math.sin(2 * math.pi * k / 8)] for k in range(8)]) * 2
def real_batch(n=256):
    return centers[torch.randint(0, 8, (n,))] + 0.05 * torch.randn(n, 2)

def mlp(out):
    return nn.Sequential(nn.Linear(2, 256), nn.LeakyReLU(0.2),
                         nn.Linear(256, 256), nn.LeakyReLU(0.2), nn.Linear(256, out))

G, D = mlp(2), mlp(1)                       # generator: noise -> point, discriminator: point -> real?
opt_g = torch.optim.Adam(G.parameters(), lr=1e-4, betas=(0.5, 0.9))
opt_d = torch.optim.Adam(D.parameters(), lr=1e-4, betas=(0.5, 0.9))
bce = nn.BCEWithLogitsLoss()

for step in range(4001):
    real, fake = real_batch(), G(torch.randn(256, 2))
    # 1) discriminator: real -> 1, fake -> 0
    d_loss = bce(D(real), torch.ones(256, 1)) + bce(D(fake.detach()), torch.zeros(256, 1))
    opt_d.zero_grad(); d_loss.backward(); opt_d.step()
    # 2) generator: fool D into saying 1
    g_loss = bce(D(fake), torch.ones(256, 1))
    opt_g.zero_grad(); g_loss.backward(); opt_g.step()
    if step % 1000 == 0:
        with torch.no_grad():
            dist, nearest = torch.cdist(G(torch.randn(4000, 2)), centers).min(1)
        hit = dist < 0.25                                  # sample landed on a real blob
        covered = (torch.bincount(nearest[hit], minlength=8) > 100).sum().item()
        print(f"step {step:4d}  D {d_loss.item():.2f}  G {g_loss.item():.2f}  "
              f"on target {hit.float().mean():.0%}  modes covered {covered}/8")