train_gan_v2.py

#!/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()