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