The denoising U-Net

pip install torch==2.5.1

import torch, torch.nn as nn, math, time

def shapes(n, size=16):
    """Tiny dataset: a white square or a white disc on black, at a random position."""
    g = torch.arange(size).float()
    yy, xx = torch.meshgrid(g, g, indexing="ij")
    c = torch.randint(4, size - 4, (n, 2, 1, 1)).float()
    dx, dy, r = xx - c[:, 0], yy - c[:, 1], 3.0
    square = ((dx.abs() < r) & (dy.abs() < r)).float()
    disc = ((dx ** 2 + dy ** 2) < r * r).float()
    pick = (torch.rand(n, 1, 1) < 0.5).float()
    return ((pick * square + (1 - pick) * disc) * 2 - 1).unsqueeze(1)   # roughly [-1, 1]

def time_embedding(t, dim=32):
    """Same sinusoidal trick transformers use for word positions, applied to the step number."""
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half).float() / half)
    a = t.float()[:, None] * freqs[None]
    return torch.cat([a.sin(), a.cos()], dim=-1)

class Block(nn.Module):
    def __init__(self, cin, cout, tdim=32):
        super().__init__()
        self.c1 = nn.Conv2d(cin, cout, 3, padding=1)
        self.c2 = nn.Conv2d(cout, cout, 3, padding=1)
        self.t  = nn.Linear(tdim, cout)                 # the step number, injected as a bias
        self.n1, self.n2 = nn.GroupNorm(4, cout), nn.GroupNorm(4, cout)
    def forward(self, x, temb):
        h = torch.nn.functional.silu(self.n1(self.c1(x)))
        h = h + self.t(temb)[:, :, None, None]
        return torch.nn.functional.silu(self.n2(self.c2(h)))

class UNet(nn.Module):
    def __init__(self, skips=True, c=32):
        super().__init__()
        self.skips = skips
        self.d1, self.d2 = Block(1, c), Block(c, c*2)
        self.mid = Block(c*2, c*2)
        self.u2 = Block(c*2 + (c*2 if skips else 0), c)
        self.u1 = Block(c + (c if skips else 0), c)
        self.out = nn.Conv2d(c, 1, 1)
        self.pool = nn.AvgPool2d(2)
        self.up = nn.Upsample(scale_factor=2, mode="nearest")
    def forward(self, x, t):
        e = time_embedding(t)
        h1 = self.d1(x, e)                       # 16x16
        h2 = self.d2(self.pool(h1), e)           # 8x8
        m  = self.mid(self.pool(h2), e)          # 4x4
        u  = self.up(m)                          # 8x8
        u  = self.u2(torch.cat([u, h2], 1) if self.skips else u, e)
        u  = self.up(u)                          # 16x16
        u  = self.u1(torch.cat([u, h1], 1) if self.skips else u, e)
        return self.out(u)

T = 200
betas = torch.linspace(1e-4, 0.02, T)
abar = torch.cumprod(1 - betas, 0)

def train(skips, steps=400, bs=64, seed=0):
    torch.manual_seed(seed)
    net = UNet(skips)
    opt = torch.optim.Adam(net.parameters(), lr=2e-3)
    last = []
    for s in range(steps):
        x0 = shapes(bs)
        t = torch.randint(0, T, (bs,))
        noise = torch.randn_like(x0)
        a = abar[t][:, None, None, None]
        xt = a.sqrt() * x0 + (1 - a).sqrt() * noise      # the forward process, in one jump
        loss = ((net(xt, t) - noise) ** 2).mean()        # predict the noise, not the image
        opt.zero_grad(); loss.backward(); opt.step()
        if s >= steps - 50: last.append(loss.item())
    return net, sum(last) / len(last)

if __name__ == "__main__":
    net = UNet(True)
    x = shapes(2)
    print("parameters:", f"{sum(p.numel() for p in net.parameters()):,}")
    print("input", tuple(x.shape), "-> output", tuple(net(x, torch.tensor([5, 199])).shape))
    print("time embedding for steps 0, 5, 199 (first 4 numbers):")
    for t in (0, 5, 199):
        print(f"  t={t:3d}  {[round(v,3) for v in time_embedding(torch.tensor([t]))[0,:4].tolist()]}")
    t0 = time.time()
    print("")
    print("skip connections   mean noise-prediction MSE over the last 50 steps")
    nets = {}
    for sk in (True, False):
        nets[sk], l = train(sk)
        print(f"  {str(sk):5s}            {l:.4f}")
    print(f"training seconds: {time.time()-t0:.0f} (varies by machine)")

    @torch.no_grad()
    def sample(net):
        torch.manual_seed(4)
        x = torch.randn(1, 1, 16, 16)
        for i in reversed(range(T)):
            e = net(x, torch.tensor([i]))
            mean = (x - betas[i] / (1 - abar[i]).sqrt() * e) / (1 - betas[i]).sqrt()
            x = mean + (betas[i].sqrt() * torch.randn_like(x) if i else 0)
        return x[0, 0]

    for sk in (True, False):
        print("")
        print("sampled from the model WITH%s skip connections:" % ("" if sk else "OUT"))
        for row in sample(nets[sk]):
            print("   " + "".join("#" if v > 0 else "." for v in row))

xt = abar[t].sqrt() * x0 + (1 - abar[t]).sqrt() * noise