Singapore team selection 2025 · Correction by IOAI Training

Image Generation with DCGAN

This page gives the full approach. Read it once you have searched.

Metric: visual quality (we add MNIST-FID)
Evaluated on: written answers
Our score: MNIST-FID 26.5, all 10 digits generated
Expected level: working DCGAN with clean digits (graded by eye, 50 pts)

Expected level reached

We solved the problem ourselves: downloaded the data, ran the experiments and measured the score below. The statement and data belong to their organisers.

The problem in brief

You must build a Deep Convolutional GAN that generates MNIST digits, filling three blanks in a notebook: the generator (Q1, 15 points), the discriminator (Q2, 15 points) and the two steps of the adversarial training loop (Q3, 20 points). The architecture must follow the DCGAN guidelines: strided convolutions instead of pooling, BatchNorm, ReLU in the generator with a Tanh output, LeakyReLU in the discriminator. The notebook fixes the rest: latent dimension 100, Adam with learning rate $2\cdot10^{-4}$ and $\beta_1 = 0.5$ , batches of 128, 15 epochs.

How you are scored

Coaches grade the code (50 points) and judge the samples by eye; the loss curves must be shown. The objective is the GAN minimax game with binary cross-entropy:

\min_G \max_D \; V(D,G) = \mathbb{E}_{x\sim p_{\text{data}}}\big[\log D(x)\big] + \mathbb{E}_{z\sim\mathcal N(0,I)}\big[\log(1 - D(G(z)))\big].

"Looks good" is not a number, so we add an objective check of our own (not part of the official grading). We train a small CNN on real MNIST (98.3 % test accuracy) and use it as a judge: on 10 000 generated digits we measure (i) the share of each predicted class, to detect mode collapse, (ii) the judge's mean confidence and (iii) a Fréchet distance between the judge's features of real test digits and of generated digits, the idea behind FID, computed with an MNIST network instead of Inception:

\text{FD} = \lVert \mu_r - \mu_g \rVert^2 + \operatorname{tr}\!\big(\Sigma_r + \Sigma_g - 2(\Sigma_r\Sigma_g)^{1/2}\big).

Two sets of real digits give FD = 2.75, so that is the floor.

First look at the data

The 60 000 MNIST digits are $28\times 28$ grey images, normalised to $[-1, 1]$ so that they live in the same range as a Tanh output. The odd size 28 is the only catch: $28 = 7\times 4$ , so the networks pass through $7\times7$ rather than a power of two. Hence the "tricky" block in Q1.

A first baseline

There is no baseline to beat here, but there is a clear failure mode to understand. We removed guideline 2 (no BatchNorm anywhere, everything else unchanged). After 5 epochs, 98 % of the samples are classified as the digit "8", FD is 1 905, and the images are faint grey blobs (right-hand panel below). That is mode collapse: the generator found one output that fools the discriminator and sends every $z$ to it. Keep this picture in mind: it shows what the guidelines are for.

The key idea

Q1, the generator: read the size formula. For a transposed convolution, output $= (\text{in} - 1)\cdot s - 2p + k$ . With the hinted settings:

block	layer	size
1	ConvT(100 to 256), $k=4, s=1, p=0$	$(1-1)\cdot1 + 4 = 4$
2	ConvT(256 to 128), $k=3, s=2, p=1$	$(4-1)\cdot2 - 2 + 3 = 7$
3	ConvT(128 to 64), $k=4, s=2, p=1$	$(7-1)\cdot2 - 2 + 4 = 14$
4	ConvT(64 to 1), $k=4, s=2, p=1$ , then Tanh	$(14-1)\cdot2 - 2 + 4 = 28$

Blocks 1 to 3 are each ConvTranspose2d, BatchNorm2d, ReLU; the last has no BatchNorm and ends in Tanh. The convolutions use bias=False because the following BatchNorm subtracts the mean anyway.

Q2, the discriminator: the mirror image. For a convolution, output $= \lfloor(\text{in} + 2p - k)/s\rfloor + 1$ : $28 \to 14 \to 7$ with $k=4, s=2, p=1$ , then a $7\times7$ convolution with no padding collapses $7\times7$ to $1\times1$ , one logit. No BatchNorm on the first layer (the DCGAN paper found that normalising the raw images destabilises training), LeakyReLU with slope 0.2 everywhere, no sigmoid at the end, because the loss is BCEWithLogitsLoss, which fuses sigmoid and log for numerical stability.

Q3, the training loop: two subtleties.

# discriminator step: real -> 1, fake -> 0
out_real = D(real_imgs);              loss_real = crit(out_real, ones)
fake = G(torch.randn(bsz, latent_dim, device=dev))
out_fake = D(fake.detach());          loss_fake = crit(out_fake, zeros)   # detach!
d_loss = loss_real + loss_fake;       d_loss.backward(); opt_d.step()
# generator step: make D say "real"
g_loss = crit(D(fake), ones);         g_loss.backward(); opt_g.step()

fake.detach() in the D step cuts the graph: the discriminator update must not push gradients into $G$ (otherwise $G$ would be trained to help $D$ ). In the G step we use fake without detach, through the freshly updated $D$ .
The generator uses real_labels: it minimises $-\log D(G(z))$ instead of the minimax term $\log(1 - D(G(z)))$ . Both have the same fixed point, but at the start, when $D$ easily rejects fakes ( $D(G(z)) \approx 0$ ), $\log(1-D)$ is flat and gives $G$ almost no gradient, while $-\log D$ is steep exactly there. This is the non-saturating loss from the original GAN paper.

Why BatchNorm saves the day. Normalising each layer's activations over the batch keeps the signal at a sane scale through the 4 layers and couples the samples of a batch, which makes it much harder for $G$ to send every $z$ to the same image. The DCGAN paper reports exactly the collapse we reproduced without it. We also use the paper's initialisation, weights drawn from $\mathcal N(0, 0.02)$ .

The full solution

solution.py contains the two networks, the DCGAN initialisation, the training loop exactly as in the notebook (15 epochs, 468 batches each), and the judge. It saves dcgan_generator_weights.pth (the file the submission asks for) and the losses. On an Apple GPU an epoch takes 40 to 90 seconds.

Samples from a fixed batch of noise vectors: after 1 epoch, noisy blobs; after 5 epochs, recognisable but broken strokes; after 15 epochs, clean digits of all classes; without BatchNorm after 5 epochs, nearly identical faint grey shapes.

Results

model	epoch	Fréchet distance (lower is better)	judge confidence	rarest digit share
DCGAN	1	494	0.65	2.7 %
DCGAN	3	103	0.82	5.7 %
DCGAN	5	114	0.83	5.2 %
DCGAN	10	48	0.87	5.5 %
DCGAN (final)	15	26.5	0.89	7.0 %
no BatchNorm (failed attempt)	5	1 905	0.25	0 % (98 % are "8")
real vs real (floor)		2.75	0.98

With 10 % per digit as perfect coverage, the final generator produces all ten digits with shares between 7 % and 14 %, and its quality keeps improving until the last epoch: more epochs would still help.

Left: moving averages of the discriminator and generator losses, which converge to about 1.0 and 1.25 and then oscillate. Middle: D(x) on real images settles near 0.63 and D(G(z)) on fakes near 0.37. Right: Fréchet distance per epoch on a log scale, falling from 494 to 26.5 for the DCGAN while the no-BatchNorm model stays near 2 000.

How to read the loss curves, the part students most often get wrong:

GAN losses do not go to zero, and they should not. They measure a game, not a fit. If $D$ were completely fooled ( $D = 0.5$ everywhere), $L_D = 2\log 2 \approx 1.39$ and $L_G = \log 2 \approx 0.69$ . Our curves settle near $L_D \approx 1.0$ and $L_G \approx 1.25$ , with $D(x) \approx 0.63$ and $D(G(z)) \approx 0.37$ : a balanced game where $D$ keeps a small edge, which is the healthy regime.
Warning signs would be $L_D \to 0$ (the discriminator wins, $G$ gets no useful gradient) or $L_G$ falling while the samples all look alike (collapse). The losses barely move between epoch 5 and 15 while the Fréchet distance drops from 114 to 26: losses alone cannot tell you about quality, so always look at samples, and ideally measure them.

What to remember

Compute every spatial size with the (transposed) convolution formulas before writing the layers; odd sizes like $7\times7$ need a non-standard kernel or padding.
detach() the fake batch in the discriminator step, and train the generator with "real" labels (the non-saturating loss) so it gets strong gradients early on.
BatchNorm is not decoration in DCGAN: without it we got mode collapse within five epochs.
GAN losses describe a balance, not progress. Judge a generator by its samples, and quantify them with a pretrained classifier: class histogram for diversity, a Fréchet distance between feature distributions for realism.

Full code (solution.py)

"""Singapore IOAI 2025 camp, CV part 1: DCGAN on MNIST (Q1 generator, Q2 discriminator, Q3 training loop).

The test is graded on visual quality, so we also MEASURE quality with a small MNIST classifier
trained on real digits ("judge"):
  * class histogram of 10 000 generated digits -> mode coverage (does G produce all ten digits?)
  * mean max-softmax confidence -> are the samples recognisable digits?
  * Frechet distance between judge features of real test digits and generated digits ("MNIST-FID")

  python solution.py judge            # train the judge (2 epochs, about 98 % test accuracy)
  python solution.py dcgan [epochs]   # train the DCGAN of the notebook, evaluate after chosen epochs
  python solution.py nobn  [epochs]   # ablation: same, without BatchNorm (guideline 2 removed)
"""
import json, sys, time
import numpy as np, torch, torch.nn as nn, torch.optim as optim
import torchvision, torchvision.transforms as T
from scipy import linalg

torch.set_num_threads(2); torch.manual_seed(0)
dev = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
tf = T.Compose([T.ToTensor(), T.Normalize((0.5,), (0.5,))])            # pixels in [-1, 1], like Tanh
train_ds = torchvision.datasets.MNIST("work/data", train=True, download=True, transform=tf)
test_ds = torchvision.datasets.MNIST("work/data", train=False, download=True, transform=tf)
latent_dim = 100

# ------------------------------------------------------------------ Q1: generator
class Generator(nn.Module):
    def __init__(self, latent_dim, ngf=64, bn=True):
        super().__init__()
        B = (lambda c: nn.BatchNorm2d(c)) if bn else (lambda c: nn.Identity())
        self.main = nn.Sequential(
            nn.ConvTranspose2d(latent_dim, ngf * 4, 4, 1, 0, bias=False), B(ngf * 4), nn.ReLU(True),  # 1 -> 4
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 3, 2, 1, bias=False), B(ngf * 2), nn.ReLU(True),     # 4 -> 7
            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False), B(ngf), nn.ReLU(True),             # 7 -> 14
            nn.ConvTranspose2d(ngf, 1, 4, 2, 1, bias=False), nn.Tanh())                               # 14 -> 28
    def forward(self, z):
        return self.main(z.view(z.size(0), -1, 1, 1))

# ------------------------------------------------------------------ Q2: discriminator
class Discriminator(nn.Module):
    def __init__(self, ndf=64, bn=True):
        super().__init__()
        B = (lambda c: nn.BatchNorm2d(c)) if bn else (lambda c: nn.Identity())
        self.main = nn.Sequential(
            nn.Conv2d(1, ndf, 4, 2, 1, bias=False), nn.LeakyReLU(0.2, True),                          # 28 -> 14
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False), B(ndf * 2), nn.LeakyReLU(0.2, True),        # 14 -> 7
            nn.Conv2d(ndf * 2, 1, 7, 1, 0, bias=False))                                               # 7 -> 1 logit
    def forward(self, img):
        return self.main(img)

def dcgan_init(m):            # DCGAN paper: all weights from N(0, 0.02), BatchNorm gamma from N(1, 0.02)
    if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): nn.init.normal_(m.weight, 0.0, 0.02)
    elif isinstance(m, nn.BatchNorm2d): nn.init.normal_(m.weight, 1.0, 0.02); nn.init.zeros_(m.bias)

# ------------------------------------------------------------------ judge (evaluation only)
class Judge(nn.Module):
    def __init__(self):
        super().__init__()
        self.f = nn.Sequential(nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
                               nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
                               nn.Flatten(), nn.Linear(64 * 49, 128), nn.ReLU())
        self.out = nn.Linear(128, 10)
    def forward(self, x, feats=False):
        h = self.f(x); return h if feats else self.out(h)

def train_judge():
    J = Judge().to(dev); opt = optim.Adam(J.parameters(), 1e-3)
    dl = torch.utils.data.DataLoader(train_ds, 128, shuffle=True)
    for ep in range(2):
        for x, y in dl:
            opt.zero_grad(); nn.functional.cross_entropy(J(x.to(dev)), y.to(dev)).backward(); opt.step()
    J.eval(); ok = 0
    with torch.no_grad():
        for x, y in torch.utils.data.DataLoader(test_ds, 1000):
            ok += (J(x.to(dev)).argmax(1).cpu() == y).sum().item()
    print("judge test accuracy", ok / len(test_ds)); torch.save(J.state_dict(), "work/judge.pt")

def frechet(a, b):
    mu1, mu2 = a.mean(0), b.mean(0); s1, s2 = np.cov(a, rowvar=False), np.cov(b, rowvar=False)
    cs = linalg.sqrtm(s1 @ s2).real
    return float(((mu1 - mu2) ** 2).sum() + np.trace(s1 + s2 - 2 * cs))

@torch.no_grad()
def evaluate(G, J, real_feats, n=10000):
    G.eval(); feats, preds, conf = [], [], []
    for k in range(0, n, 1000):
        x = G(torch.randn(1000, latent_dim, device=dev))
        f = J(x, feats=True); p = J.out(f).softmax(1)
        feats.append(f.cpu().numpy()); preds.append(p.argmax(1).cpu().numpy()); conf.append(p.max(1).values.cpu().numpy())
    G.train()
    feats, preds = np.concatenate(feats), np.concatenate(preds)
    hist = np.bincount(preds, minlength=10) / n
    return {"fid": frechet(real_feats, feats), "confidence": float(np.concatenate(conf).mean()),
            "class_hist": hist.round(4).tolist(), "min_class_share": float(hist.min())}

# ------------------------------------------------------------------ Q3: training loop
def train_dcgan(num_epochs, bn=True, tag="dcgan", eval_at=(1, 3, 5, 10, 15)):
    J = Judge().to(dev); J.load_state_dict(torch.load("work/judge.pt")); J.eval()
    with torch.no_grad():
        real = np.concatenate([J(x.to(dev), feats=True).cpu().numpy() for x, _ in torch.utils.data.DataLoader(test_ds, 1000)])
    print("FID real train vs real test (floor):", round(frechet(real, np.concatenate(
        [J(x.to(dev), feats=True).detach().cpu().numpy() for x, _ in torch.utils.data.DataLoader(torch.utils.data.Subset(train_ds, range(10000)), 1000)])), 3))
    G, D = Generator(latent_dim, bn=bn).to(dev), Discriminator(bn=bn).to(dev)
    G.apply(dcgan_init); D.apply(dcgan_init)
    crit = nn.BCEWithLogitsLoss()
    opt_g = optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))
    opt_d = optim.Adam(D.parameters(), lr=2e-4, betas=(0.5, 0.999))
    loader = torch.utils.data.DataLoader(train_ds, batch_size=128, shuffle=True, drop_last=True)
    fixed = torch.randn(64, latent_dim, device=dev)
    d_losses, g_losses, log, D_real, D_fake = [], [], [], [], []
    for epoch in range(1, num_epochs + 1):
        t0 = time.time()
        for real_imgs, _ in loader:
            real_imgs = real_imgs.to(dev); bsz = real_imgs.size(0)
            ones = torch.ones(bsz, 1, 1, 1, device=dev); zeros = torch.zeros(bsz, 1, 1, 1, device=dev)
            # --- Q3.1 discriminator: push D(x) -> 1 on real, D(G(z)) -> 0 on fake
            opt_d.zero_grad()
            out_real = D(real_imgs); loss_real = crit(out_real, ones)
            fake = G(torch.randn(bsz, latent_dim, device=dev))
            out_fake = D(fake.detach()); loss_fake = crit(out_fake, zeros)     # detach: no gradient into G
            d_loss = loss_real + loss_fake; d_loss.backward(); opt_d.step()
            # --- Q3.2 generator: non-saturating loss, push D(G(z)) -> 1 ("real" labels)
            opt_g.zero_grad()
            g_loss = crit(D(fake), ones); g_loss.backward(); opt_g.step()
            d_losses.append(d_loss.item()); g_losses.append(g_loss.item())
            D_real.append(torch.sigmoid(out_real).mean().item()); D_fake.append(torch.sigmoid(out_fake).mean().item())
        msg = {"epoch": epoch, "sec": round(time.time() - t0), "d_loss": float(np.mean(d_losses[-468:])),
               "g_loss": float(np.mean(g_losses[-468:]))}
        if epoch in eval_at:
            msg.update(evaluate(G, J, real))
            G.eval()
            with torch.no_grad(): torch.save(G(fixed).cpu(), f"work/{tag}_grid_ep{epoch}.pt")
            G.train()
        print(msg, flush=True); log.append(msg)
    torch.save(G.state_dict(), f"work/{tag}_generator_weights.pth")
    json.dump({"log": log, "d_losses": d_losses, "g_losses": g_losses, "D_real": D_real, "D_fake": D_fake},
              open(f"work/{tag}_log.json", "w"))

if __name__ == "__main__":
    what = sys.argv[1]; ep = int(sys.argv[2]) if len(sys.argv) > 2 else 15
    if what == "judge": train_judge()
    if what == "dcgan": train_dcgan(ep)
    if what == "nobn": train_dcgan(ep, bn=False, tag="nobn", eval_at=(1, 3, 5))