Add training script

train_gan_v2.py

@@ -0,0 +1,154 @@
+#!/usr/bin/env python3
+"""Small DCGAN for 64x64 RGB company logos. GPU-enabled.
+Loads /work/logos/logos64.npy (N,3,64,64) float32 in [0,1].
+Usage: python3 train_gan_gpu.py [--steps 12000] [--batch 128] [--seed 7]
+"""
+import os, sys, argparse, time, random
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+from PIL import Image
+
+OUT = "gan_out"; os.makedirs(OUT, exist_ok=True)
+LATENT = 100
+CH, SZ = 3, 64
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def parse():
+    p = argparse.ArgumentParser()
+    p.add_argument("--steps", type=int, default=12000)
+    p.add_argument("--batch", type=int, default=128)
+    p.add_argument("--seed", type=int, default=7)
+    p.add_argument("--ckpt_every", type=int, default=2000)
+    return p.parse_args()
+
+class G(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc = nn.Linear(LATENT, 512 * 8 * 8)
+        self.body = nn.Sequential(
+            nn.BatchNorm2d(512), nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(512, 256, 4, 2, 1),
+            nn.BatchNorm2d(256), nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(256, 128, 4, 2, 1),
+            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
+            nn.ConvTranspose2d(128, 64, 4, 2, 1),
+            nn.BatchNorm2d(64), nn.ReLU(inplace=True),
+            nn.Conv2d(64, CH, 3, 1, 1),
+            nn.Tanh(),
+        )
+    def forward(self, z):
+        x = self.fc(z).view(-1, 512, 8, 8)
+        return self.body(x)
+
+class D(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.body = nn.Sequential(
+            nn.Conv2d(CH, 64, 4, 2, 1), nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(64, 128, 4, 2, 1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(128, 256, 4, 2, 1), nn.BatchNorm2d(256), nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(256, 512, 4, 2, 1), nn.BatchNorm2d(512), nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(512, 1, 4, 1),
+        )
+    def forward(self, x):
+        return self.body(x).view(-1, 1)
+
+def make_grid(imgs, ncols=8):
+    n = imgs.shape[0]
+    rows = (n + ncols - 1) // ncols
+    pad = np.zeros((rows * ncols, 3, SZ, SZ), dtype=np.float32) + 1.0
+    for i in range(n):
+        pad[i] = imgs[i]
+    tiles = []
+    for r in range(rows):
+        row = []
+        for c in range(ncols):
+            idx = r * ncols + c
+            if idx < n:
+                arr = (pad[idx] * 0.5 + 0.5).clip(0, 1)
+                row.append((arr.transpose(1, 2, 0) * 255).astype(np.uint8))
+            else:
+                row.append(np.full((SZ, SZ, 3), 255, dtype=np.uint8))
+        tiles.append(np.concatenate(row, axis=1))
+    return np.concatenate(tiles, axis=0)
+
+def main():
+    a = parse()
+    torch.manual_seed(a.seed); np.random.seed(a.seed); random.seed(a.seed)
+    X = np.load("logos/logos64.npy")
+    print(f"device={DEVICE}")
+    print(f"data {X.shape} min={X.min():.3f} max={X.max():.3f} mean={X.mean():.3f}")
+    X = (X * 2.0 - 1.0).astype(np.float32)
+    ds = TensorDataset(torch.from_numpy(X))
+    dl = DataLoader(ds, batch_size=a.batch, shuffle=True, drop_last=True, num_workers=0)
+
+    gen = G().to(DEVICE); disc = D().to(DEVICE)
+    ng = sum(p.numel() for p in gen.parameters())
+    nd = sum(p.numel() for p in disc.parameters())
+    print(f"generator params {ng:,}  discriminator params {nd:,}  total {ng+nd:,}")
+    with open(os.path.join(OUT, "param_count.txt"), "w") as f:
+        f.write(f"generator={ng}\ndiscriminator={nd}\ntotal={ng+nd}\n")
+
+    opt_g = torch.optim.Adam(gen.parameters(), lr=2e-4, betas=(0.5, 0.999))
+    opt_d = torch.optim.Adam(disc.parameters(), lr=4e-5, betas=(0.5, 0.999))
+    bce = nn.BCEWithLogitsLoss()
+
+    fixed_z = torch.randn(64, LATENT, device=DEVICE)
+    step = 0
+    t0 = time.time()
+    for epoch in range(10_000):
+        for xb in dl:
+            xb = xb[0].to(DEVICE, non_blocking=True)
+            bs = xb.size(0)
+            real = xb
+            z = torch.randn(bs, LATENT, device=DEVICE)
+            fake = gen(z)
+            opt_d.zero_grad()
+            d_real = disc(real)
+            d_fake = disc(fake.detach())
+            # label smoothing: real=0.9, fake=0.1
+            real_t = torch.full_like(d_real, 0.9)
+            fake_t = torch.full_like(d_fake, 0.1)
+            loss_d = (bce(d_real, real_t) + bce(d_fake, fake_t)).mean()
+            # R1 gradient penalty (regularize D, prevents collapse)
+            inp = real.clone().requires_grad_(True)
+            out = disc(inp)
+            grads = torch.autograd.grad(outputs=out, inputs=inp,
+                       grad_outputs=torch.ones_like(out),
+                       create_graph=True, retain_graph=True)[0]
+            pen = (grads**2).sum(dim=[1,2,3]).mean()
+            loss_d = loss_d + 10.0 * pen
+            loss_d.backward(); opt_d.step()
+            opt_g.zero_grad()
+            loss_g = bce(disc(gen(z)), torch.ones(bs, 1, device=DEVICE)).mean()
+            loss_g.backward(); opt_g.step()
+            step += 1
+            if step % 200 == 0:
+                el = (time.time() - t0) / step
+                print(f"step {step}  d={loss_d.item():.4f}  g={loss_g.item():.4f}  {el*1000:.0f}ms/step  eta {el*(a.steps-step)/60:.1f}m", flush=True)
+            if step % a.ckpt_every == 0:
+                torch.save({"G": gen.state_dict(), "D": disc.state_dict(), "step": step},
+                           os.path.join(OUT, f"ckpt_{step}.pt"))
+                if step >= a.ckpt_every:
+                    gen.eval()
+                    with torch.no_grad():
+                        samp = gen(fixed_z).cpu().numpy()
+                    grid = make_grid(samp)
+                    Image.fromarray(grid).save(os.path.join(OUT, f"samples_{step}.png"))
+                    gen.train()
+                    print(f"  saved ckpt_{step}.pt + samples_{step}.png", flush=True)
+            if step >= a.steps:
+                break
+        if step >= a.steps:
+            break
+    torch.save({"G": gen.state_dict(), "D": disc.state_dict(), "step": step}, os.path.join(OUT, "final.pt"))
+    gen.eval()
+    with torch.no_grad():
+        samp = gen(fixed_z).cpu().numpy()
+    Image.fromarray(make_grid(samp)).save(os.path.join(OUT, "samples_final.png"))
+    print(f"DONE step={step}  elapsed {(time.time()-t0)/60:.1f}m", flush=True)
+
+if __name__ == "__main__":
+    main()