train_diffusion.py

"""Proof-of-life DDPM on CIFAR-10.

Trains a small U-Net to predict noise (Ho et al. 2020, simplified).
Saves sample grids and checkpoints to diffusion/out/ as it goes.

Run:  python diffusion/train_diffusion.py [--base 128 --epochs 300 --out out_big --resume]
"""
import argparse
import copy
import math
import os
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import save_image

# ---------------- config ----------------
DATA_ROOT = os.path.join(os.path.dirname(__file__), "..", "cifar-10-python")
BATCH_SIZE = 128
LR = 2e-4
T = 1000                 # diffusion timesteps
DDIM_STEPS = 50          # sampling steps
DEVICE = "cuda"

# ---------------- model ----------------
class TimeEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.mlp = nn.Sequential(nn.Linear(dim, dim * 4), nn.SiLU(), nn.Linear(dim * 4, dim * 4))

    def forward(self, t):
        half = self.dim // 2
        freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
        args = t[:, None].float() * freqs[None]
        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
        return self.mlp(emb)


class ResBlock(nn.Module):
    def __init__(self, in_ch, out_ch, emb_dim):
        super().__init__()
        self.norm1 = nn.GroupNorm(8, in_ch)
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.emb_proj = nn.Linear(emb_dim, out_ch)
        self.norm2 = nn.GroupNorm(8, out_ch)
        self.dropout = nn.Dropout(0.1)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()

    def forward(self, x, emb):
        h = self.conv1(F.silu(self.norm1(x)))
        h = h + self.emb_proj(F.silu(emb))[:, :, None, None]
        h = self.conv2(self.dropout(F.silu(self.norm2(h))))
        return h + self.skip(x)


class Attention(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.norm = nn.GroupNorm(8, ch)
        self.qkv = nn.Conv2d(ch, ch * 3, 1)
        self.proj = nn.Conv2d(ch, ch, 1)

    def forward(self, x):
        b, c, h, w = x.shape
        q, k, v = self.qkv(self.norm(x)).reshape(b, 3, c, h * w).unbind(1)
        attn = torch.softmax(q.transpose(1, 2) @ k / math.sqrt(c), dim=-1)
        out = (v @ attn.transpose(1, 2)).reshape(b, c, h, w)
        return x + self.proj(out)


class UNet(nn.Module):
    """Small U-Net: 32 -> 16 -> 8, attention at 16 and 8."""

    def __init__(self, base=64):
        super().__init__()
        chs = [base, base * 2, base * 2]  # per resolution level
        emb_dim = base * 4
        self.time_emb = TimeEmbedding(base)
        self.stem = nn.Conv2d(3, base, 3, padding=1)

        # down
        self.d1a = ResBlock(base, chs[0], emb_dim)
        self.d1b = ResBlock(chs[0], chs[0], emb_dim)
        self.down1 = nn.Conv2d(chs[0], chs[0], 3, stride=2, padding=1)
        self.d2a = ResBlock(chs[0], chs[1], emb_dim)
        self.d2b = ResBlock(chs[1], chs[1], emb_dim)
        self.attn2 = Attention(chs[1])
        self.down2 = nn.Conv2d(chs[1], chs[1], 3, stride=2, padding=1)
        self.d3a = ResBlock(chs[1], chs[2], emb_dim)
        self.d3b = ResBlock(chs[2], chs[2], emb_dim)
        self.attn3 = Attention(chs[2])

        # middle
        self.mid1 = ResBlock(chs[2], chs[2], emb_dim)
        self.mid_attn = Attention(chs[2])
        self.mid2 = ResBlock(chs[2], chs[2], emb_dim)

        # up
        self.u3a = ResBlock(chs[2] * 2, chs[2], emb_dim)
        self.u3b = ResBlock(chs[2], chs[1], emb_dim)
        self.uattn3 = Attention(chs[1])
        self.up2 = nn.ConvTranspose2d(chs[1], chs[1], 4, stride=2, padding=1)
        self.u2a = ResBlock(chs[1] * 2, chs[1], emb_dim)
        self.u2b = ResBlock(chs[1], chs[0], emb_dim)
        self.uattn2 = Attention(chs[0])
        self.up1 = nn.ConvTranspose2d(chs[0], chs[0], 4, stride=2, padding=1)
        self.u1a = ResBlock(chs[0] * 2, chs[0], emb_dim)
        self.u1b = ResBlock(chs[0], chs[0], emb_dim)

        self.out_norm = nn.GroupNorm(8, chs[0])
        self.out = nn.Conv2d(chs[0], 3, 3, padding=1)

    def forward(self, x, t):
        emb = self.time_emb(t)
        h1 = self.stem(x)
        h1 = self.d1b(self.d1a(h1, emb), emb)                    # 32
        h2 = self.down1(h1)
        h2 = self.attn2(self.d2b(self.d2a(h2, emb), emb))        # 16
        h3 = self.down2(h2)
        h3 = self.attn3(self.d3b(self.d3a(h3, emb), emb))        # 8

        m = self.mid2(self.mid_attn(self.mid1(h3, emb)), emb)

        u = self.u3b(self.u3a(torch.cat([m, h3], 1), emb), emb)
        u = self.uattn3(u)
        u = self.up2(u)
        u = self.u2b(self.u2a(torch.cat([u, h2], 1), emb), emb)
        u = self.uattn2(u)
        u = self.up1(u)
        u = self.u1b(self.u1a(torch.cat([u, h1], 1), emb), emb)
        return self.out(F.silu(self.out_norm(u)))


# ---------------- diffusion ----------------
class Diffusion:
    def __init__(self, T, device):
        self.T = T
        beta = torch.linspace(1e-4, 0.02, T, device=device)
        self.alpha_bar = torch.cumprod(1 - beta, dim=0)

    def add_noise(self, x0, t, noise):
        ab = self.alpha_bar[t][:, None, None, None]
        return ab.sqrt() * x0 + (1 - ab).sqrt() * noise

    @torch.no_grad()
    def ddim_sample(self, model, n, steps, device):
        ts = torch.linspace(self.T - 1, 0, steps, device=device).long()
        x = torch.randn(n, 3, 32, 32, device=device)
        for i in range(steps):
            t = ts[i].expand(n)
            eps = model(x, t)
            ab_t = self.alpha_bar[ts[i]]
            x0 = (x - (1 - ab_t).sqrt() * eps) / ab_t.sqrt()
            x0 = x0.clamp(-1, 1)
            if i + 1 < steps:
                ab_prev = self.alpha_bar[ts[i + 1]]
                x = ab_prev.sqrt() * x0 + (1 - ab_prev).sqrt() * eps
            else:
                x = x0
        return x


# ---------------- training ----------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", type=int, default=64, help="U-Net base channels")
    ap.add_argument("--epochs", type=int, default=50)
    ap.add_argument("--out", default="out", help="output dir name (under diffusion/)")
    ap.add_argument("--sample-every", type=int, default=5)
    ap.add_argument("--ema-decay", type=float, default=0.999)
    ap.add_argument("--resume", action="store_true", help="resume from checkpoint.pt in out dir")
    args = ap.parse_args()

    out_dir = os.path.join(os.path.dirname(__file__), args.out)
    os.makedirs(out_dir, exist_ok=True)
    torch.backends.cudnn.benchmark = True

    tfm = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.5] * 3, [0.5] * 3),  # -> [-1, 1]
    ])
    ds = datasets.CIFAR10(DATA_ROOT, train=True, transform=tfm, download=True)
    dl = DataLoader(ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=4,
                    pin_memory=True, persistent_workers=True, drop_last=True)

    model = UNet(base=args.base).to(DEVICE)
    ema = copy.deepcopy(model).requires_grad_(False)
    opt = torch.optim.Adam(model.parameters(), lr=LR)
    scaler = torch.amp.GradScaler()
    diff = Diffusion(T, DEVICE)

    start_epoch = 1
    ckpt_path = os.path.join(out_dir, "checkpoint.pt")
    if args.resume and os.path.exists(ckpt_path):
        ck = torch.load(ckpt_path, map_location=DEVICE, weights_only=True)
        model.load_state_dict(ck["model"])
        ema.load_state_dict(ck["ema"])
        opt.load_state_dict(ck["opt"])
        start_epoch = ck["epoch"] + 1
        print(f"resumed from epoch {ck['epoch']}", flush=True)

    n_params = sum(p.numel() for p in model.parameters())
    print(f"params: {n_params/1e6:.1f}M | batches/epoch: {len(dl)}", flush=True)

    step = 0
    for epoch in range(start_epoch, args.epochs + 1):
        t0 = time.time()
        running = 0.0
        for x, _ in dl:
            x = x.to(DEVICE, non_blocking=True)
            t = torch.randint(0, T, (x.size(0),), device=DEVICE)
            noise = torch.randn_like(x)
            xt = diff.add_noise(x, t, noise)

            with torch.amp.autocast("cuda"):
                loss = F.mse_loss(model(xt, t), noise)
            opt.zero_grad(set_to_none=True)
            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()

            with torch.no_grad():
                for pe, pm in zip(ema.parameters(), model.parameters()):
                    pe.lerp_(pm, 1 - args.ema_decay)
                for be, bm in zip(ema.buffers(), model.buffers()):
                    be.copy_(bm)

            running += loss.item()
            step += 1

        dt = time.time() - t0
        print(f"epoch {epoch:3d}/{args.epochs} | loss {running/len(dl):.4f} | "
              f"{dt:.0f}s | {len(dl)/dt:.1f} it/s", flush=True)

        if epoch % args.sample_every == 0 or epoch == args.epochs:
            ema.eval()
            imgs = diff.ddim_sample(ema, 64, DDIM_STEPS, DEVICE)
            save_image(imgs * 0.5 + 0.5, os.path.join(out_dir, f"samples_epoch{epoch:03d}.png"), nrow=8)
            torch.save({"model": model.state_dict(), "ema": ema.state_dict(),
                        "opt": opt.state_dict(), "epoch": epoch},
                       ckpt_path)
            print(f"  -> saved samples_epoch{epoch:03d}.png + checkpoint", flush=True)


if __name__ == "__main__":
    main()