| import os, sys, time, math, argparse |
| import numpy as np, torch, torch.nn as nn |
| from torch.utils.data import DataLoader, TensorDataset |
| import PIL.Image as Image |
|
|
| torch.manual_seed(0); np.random.seed(0) |
|
|
| A = argparse.ArgumentParser() |
| A.add_argument("--data", default="/work/logos/logos64.npy") |
| A.add_argument("--out", default="/work/models/logo-gan") |
| A.add_argument("--steps", type=int, default=4000) |
| A.add_argument("--batch", type=int, default=64) |
| A.add_argument("--zdim", type=int, default=128) |
| A.add_argument("--lr", type=float, default=2e-4) |
| A.add_argument("--ckpt-every", type=int, default=500) |
| A.add_argument("--sample-every", type=int, default=500) |
| A.add_argument("--device", default="cuda") |
| a = A.parse_args() |
|
|
| os.makedirs(a.out, exist_ok=True) |
| log = open(os.path.join(a.out, "train.log"), "w") |
| def L(*x): |
| s = " ".join(str(v) for v in x) |
| print(s, flush=True); log.write(s + "\n"); log.flush() |
|
|
| dev = a.device if (a.device == "cuda" and torch.cuda.is_available()) else "cpu" |
| L("device", dev) |
|
|
| data = np.load(a.data) |
| L("data", data.shape, "range", float(data.min()), float(data.max())) |
| X = torch.from_numpy(data).permute(0,3,1,2).to(dev) |
| ds = TensorDataset(X); dl = DataLoader(ds, batch_size=a.batch, shuffle=True, drop_last=True) |
| n = len(data) |
|
|
| def conv_block(in_c, out_c, ks=4, s=2, p=1): |
| return nn.Sequential( |
| nn.Conv2d(in_c, out_c, ks, s, p, bias=False), |
| nn.BatchNorm2d(out_c), nn.LeakyReLU(0.2, inplace=True)) |
|
|
| class G(nn.Module): |
| def __init__(self, zdim=128): |
| super().__init__() |
| self.fc = nn.Sequential(nn.Linear(zdim, 256*8*8), nn.BatchNorm1d(256*8*8), nn.ReLU()) |
| self.up = nn.Sequential( |
| nn.ConvTranspose2d(256,128,4,2,1,bias=False), nn.BatchNorm2d(128), nn.ReLU(), |
| nn.ConvTranspose2d(128,64,4,2,1,bias=False), nn.BatchNorm2d(64), nn.ReLU(), |
| nn.ConvTranspose2d(64,3,4,2,1,bias=True), nn.Tanh()) |
| def forward(self, z): |
| h = self.fc(z).view(-1,256,8,8) |
| return self.up(h) |
|
|
| class D(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Conv2d(3,64,4,2,1,bias=False), nn.BatchNorm2d(64), nn.LeakyReLU(0.2,inplace=True), |
| nn.Conv2d(64,128,4,2,1,bias=False), nn.BatchNorm2d(128), nn.LeakyReLU(0.2,inplace=True), |
| nn.Conv2d(128,256,4,2,1,bias=False), nn.BatchNorm2d(256), nn.LeakyReLU(0.2,inplace=True), |
| nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(256,1)) |
| def forward(self, x): return self.net(x).squeeze(-1) |
|
|
| g = G(a.zdim).to(dev); d = D().to(dev) |
| opt_g = torch.optim.Adam(g.parameters(), a.lr, betas=(0.5,0.999)) |
| opt_d = torch.optim.Adam(d.parameters(), a.lr, betas=(0.5,0.999)) |
| BCE = nn.BCEWithLogitsLoss() |
| zdim = a.zdim |
|
|
| def sample_grid(zs, path, step): |
| with torch.no_grad(): |
| imgs = g(zs).cpu().numpy() |
| imgs = (imgs*0.5+0.5).clip(0,1).transpose(0,2,3,1) |
| imgs = (imgs*255).astype(np.uint8) |
| |
| cols = rows = 8 |
| cell = 64; gap = 2 |
| W = cols*cell + (cols+1)*gap; H = rows*cell + (rows+1)*gap |
| canvas = np.full((H, W, 3), 255, dtype=np.uint8) |
| for i in range(min(len(imgs), cols*rows)): |
| r = i // cols; c = i % cols |
| y = gap + r*(cell+gap); x = gap + c*(cell+gap) |
| canvas[y:y+cell, x:x+cell] = imgs[i] |
| Image.fromarray(canvas).save(path) |
|
|
| def gen_samples(step): |
| zs = torch.randn(a.batch, zdim, device=dev) |
| name = f"grid_{step:05d}.png" if isinstance(step,int) else f"grid_{step}.png" |
| sample_grid(zs, os.path.join(a.out, name), step) |
|
|
| it = iter(dl) |
| best_g = None |
| for step in range(1, a.steps+1): |
| try: |
| real = next(it) |
| except StopIteration: |
| it = iter(dl); real = next(it) |
| real = real if isinstance(real, torch.Tensor) else real[0] |
| b = real.size(0) |
| |
| opt_d.zero_grad() |
| noise = torch.randn(b, zdim, device=dev) |
| fake = g(noise).detach() |
| loss_d = BCE(d(real), torch.ones(b, device=dev)) + BCE(d(fake), torch.zeros(b, device=dev)) |
| loss_d.backward(); opt_d.step() |
| |
| opt_g.zero_grad() |
| noise = torch.randn(b, zdim, device=dev) |
| fake = g(noise) |
| loss_g = BCE(d(fake), torch.ones(b, device=dev)) |
| loss_g.backward(); opt_g.step() |
| if step % 100 == 0: |
| L(f"step {step}/{a.steps} d_loss {loss_d.item():.4f} g_loss {loss_g.item():.4f}") |
| if step % a.ckpt_every == 0: |
| torch.save({"g": g.state_dict(), "d": d.state_dict(), "step": step, "zdim": zdim}, |
| os.path.join(a.out, "ckpt.pt")) |
| gen_samples(step) |
| if step % a.sample_every == 0: |
| gen_samples(step) |
|
|
| |
| torch.save({"g": g.state_dict(), "d": d.state_dict(), "step": a.steps, "zdim": zdim}, |
| os.path.join(a.out, "final.pt")) |
| gen_samples(a.steps) |
| |
| torch.manual_seed(1234) |
| gen_samples("final") |
| L(f"DONE final step {a.steps}") |
| L("ckpt files:", sorted(os.listdir(a.out))) |
| log.close() |
|
|
