ImageGen

GANs and Diffusion behind One Training Loop

Generative AI · Jun 2024 · In progress

DCGAN, a ResNet GAN and a diffusion U-Net that train, checkpoint and sample the same way, from MNIST on a CPU to animal faces on GPUs.

3 architectures
Multi-GPU with torchrun
AFHQ, 16,130 images

Built withPyTorchtorchvisionHugging Face Datasetstorchrunpytest

View code on GitHub

Overview

I wanted to generate animal faces and to understand both families of models that can do it: GANs, where a generator and a discriminator are trained against each other, and diffusion, where a U-Net learns to remove a little noise and generation is denoising pure noise step by step. ImageGen puts a DCGAN, a SNGAN/BigGAN-style ResNet GAN and a DDPM U-Net behind one training loop, one sampler and one checkpoint format in plain PyTorch.

All three can be class-conditional ('draw a cat'), keep an EMA copy of the weights for sampling, resume from an atomic checkpoint that stores optimisers, EMA, RNG and configs, and run on a CPU, one GPU, several GPUs or several machines through torchrun. Presets cover MNIST, CIFAR-10 and AFHQ (16,130 animal faces in three classes: cat, dog, wild) at 64 and 128 px, with a prepare command that streams AFHQ from the Hugging Face Hub into resized class folders. FID can be computed during training and best.pt follows the lowest score.

The project started as a Colab notebook with a ResBlock GAN on a zip of animal photos; it trained on the same 32 noise vectors every step and had several other bugs, all listed and fixed in the package. So far every architecture and GAN loss has been checked on MNIST on a CPU, including the combinations that did not work. The full GPU runs on CIFAR-10 and AFHQ, FID numbers and published checkpoints are the next step.

The question

Can one codebase train both GAN families and a diffusion model with the same loop, checkpoint and sampling interface, and scale from a CPU sanity run on MNIST to AFHQ animal faces on several GPUs?

Key features

Three architectures behind one --arch flag: DCGAN for any power-of-two resolution, a ResNet GAN with spectral norm, hinge loss, conditional BatchNorm, projection discriminator and self-attention, and a diffusion U-Net with cosine schedule, ε or v prediction, DDIM sampling and classifier-free guidance.
Four GAN losses (non-saturating, hinge, WGAN-GP, LSGAN), the R1 penalty and DiffAugment for small datasets, selectable by flag.
One task interface, GANTask / DiffusionTask, so the training loop, checkpointing, EMA and sampling do not know which family they run.
Named presets for MNIST, CIFAR-10, AFHQ 64 / 128 and an offline shapes set, with every hyperparameter overridable from the command line.
CPU, single-GPU, multi-GPU and multi-node training with bf16 autocast on recent GPUs, a time budget that stops and saves cleanly, and exact resume.
imagegen-sample for grids, per-class generation and latent walks, imagegen-fid for Inception FID, and imagegen-prepare for AFHQ, folders or zips of photos.
A Kaggle/Colab notebook built from the package source, with a live dashboard, a 12-hour budget and resume.

How it works

Data in [-1, 1]
Images are resized on the short side, centre-cropped and scaled to [-1, 1], so the generators' tanh output and the diffusion noise share the data scale. Sub-folders become classes.
A GAN step
d_steps times: make fakes without gradients into G, score real and fake, add a gradient penalty or R1 if asked, update D. Then freeze D, make fresh fakes, update G through D, and update the EMA copy of G used for sampling.
A diffusion step
Add noise at a random timestep of the cosine schedule and regress ε or v with a mean squared error; sampling runs DDIM (or DDPM-style with η = 1) from pure noise, with classifier-free guidance for conditional models.
One loop for both
The loop only calls task.train_step(), writes last.pt and preview grids from fixed noise every few steps, logs metrics to a JSONL file and, when asked, computes FID and keeps best.pt.
Scale out
Under torchrun each process computes its own gradients and one all-reduce averages them before every optimiser step; only rank 0 writes files. A test runs two CPU processes and checks they end with identical weights.

Examples

Conditional ResNet GAN (hinge loss, spectral norm, self-attention, projection discriminator, 32 channels, 2 D steps) after 3,000 CPU steps: the best GAN result, each row is one digit class and the rows already followed the labels at step 500.

Conditional DCGAN with the non-saturating loss after 4,000 CPU steps: the correct digit per row from step 500, with d_loss around 0.3 and g_loss around 3 throughout.

Diffusion U-Net with v-prediction after 2,500 steps (16 channels, guidance 2): clean black backgrounds and the correct digit per row; the loss went from 1.0 to 0.095.

The same U-Net with ε-prediction: the digits are right but many samples keep grey, speckled backgrounds, because recovering x0 from ε amplifies errors about 70x at the noisiest steps. This is why the presets switched to v-prediction.

Input

imagegen-train --preset mnist-dcgan --conditional --steps 4000 --output-dir checkpoints/mnist-dcgan imagegen-sample --checkpoint checkpoints/mnist-dcgan/last.pt --class 7 --num 16 --out sevens.png

Output

checkpoints/mnist-dcgan/last.pt, samples/step_*.png and metrics.jsonl, then sevens.png: 16 generated sevens from the EMA generator

The quick-start pair: train a conditional DCGAN on MNIST, then ask the checkpoint for one class. The checkpoint stores its configs, so sampling needs no architecture flags.

Tech stack

PyTorch
Models, losses, the diffusion process, the training loop and distributed training
torchvision
MNIST, Fashion-MNIST and CIFAR-10, and the Inception-v3 features for FID
Hugging Face Datasets
Streams AFHQ from the Hub into resized class folders
torchrun
Multi-GPU and multi-node launches with one all-reduce per step
pytest
Offline CPU tests for every module, including a two-process distributed check

What I learned

ε-prediction amplifies denoiser errors about 70x at the noisiest steps; v-prediction fixed the backgrounds and became the default
A per-sample gradient penalty and BatchNorm in the critic do not mix; spectral norm does
The first notebook trained the generator on the same 32 noise vectors every step, so it memorised 32 images; fresh noise per step is the fix

Limitations

Only CPU sanity runs on MNIST are measured; the CIFAR-10 and AFHQ GPU runs, FID numbers and published checkpoints are still on the roadmap

Keywords

image generation · GAN · DCGAN · SNGAN · BigGAN · spectral normalization · hinge loss · projection discriminator · DiffAugment · diffusion models · DDPM · DDIM · v-prediction · classifier-free guidance · EMA · AFHQ · CIFAR-10 · MNIST · distributed training · PyTorch