Add training script

train_logo_gan_v2.py

@@ -0,0 +1,93 @@
+import os, sys, time, math, argparse
+import numpy as np, torch, torch.nn as nn
+import PIL.Image as Image
+
+torch.manual_seed(0); np.random.seed(0)
+
+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)
+
+def grid(imgs, path):
+    imgs=(imgs*0.5+0.5).clip(0,1).transpose(0,2,3,1).astype(np.uint8) if imgs.min()<-0.5 else imgs.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 main():
+    a=argparse.Namespace(data="/work/logos/logos_big.npy", out="/work/models/logo-gan-v2",
+                         steps=8000, bs=16, zdim=128, lr=2e-4, betas=(0.5,0.999), dev="cuda",
+                         ckpt_every=1000, sample_every=1000)
+    os.makedirs(a.out, exist_ok=True)
+    dev = a.dev if (a.dev=="cuda" and torch.cuda.is_available()) else "cpu"
+    imgs = np.load(a.data).astype(np.float32)          # (N,64,64,3) 0..1
+    imgs = (imgs - 0.5)/0.5                            # -> [-1,1]
+    X = torch.from_numpy(imgs).permute(0,3,1,2).to(dev)
+    N = X.shape[0]
+    g = G(a.zdim).to(dev); d = D().to(dev)
+    go = torch.optim.Adam(g.parameters(), a.lr, betas=a.betas)
+    do = torch.optim.Adam(d.parameters(), a.lr, betas=a.betas)
+    logf = open(os.path.join(a.out,"train.log"),"a")
+    def lg(s): print(s, flush=True); logf.write(s+"\n"); logf.flush()
+    lg(f"START v2 N={N} dev={dev} steps={a.steps} bs={a.bs} zdim={a.zdim} lr={a.lr}")
+    g_params=sum(p.numel() for p in g.parameters()); d_params=sum(p.numel() for p in d.parameters())
+    lg(f"params generator={g_params} discriminator={d_params} total={g_params+d_params}")
+    torch.manual_seed(0); np.random.seed(0)
+    for step in range(1, a.steps+1):
+        idx = torch.randint(0, N, (a.bs,))
+        real = X[idx]
+        z = torch.randn(a.bs, a.zdim, device=dev)
+        # D
+        do.zero_grad()
+        fake = g(z).detach()
+        dloss = -(torch.log(torch.sigmoid(d(real)).clamp(min=1e-7)).mean() + torch.log(1-torch.sigmoid(d(fake)).clamp(min=1e-7)).mean())
+        dloss.backward(); do.step()
+        # G
+        go.zero_grad()
+        fake = g(z)
+        gloss = -torch.log(torch.sigmoid(d(fake)).clamp(min=1e-7)).mean()
+        gloss.backward(); go.step()
+        if step % 200 == 0:
+            lg(f"step {step}/{a.steps} dloss={float(dloss):.4f} gloss={float(gloss):.4f}")
+        if step % a.ckpt_every == 0:
+            torch.save({"g":g.state_dict(),"d":d.state_dict(),"step":step,"zdim":a.zdim}, os.path.join(a.out,f"ckpt_{step:05d}.pt"))
+        if step % a.sample_every == 0:
+            g.eval()
+            with torch.no_grad():
+                fake = g(torch.randn(64, a.zdim, device=dev)).cpu().numpy()
+            grid(fake, os.path.join(a.out, f"grid_{step:05d}.png"))
+            g.train()
+    # final
+    g.eval()
+    with torch.no_grad():
+        fake = g(torch.randn(64, a.zdim, device=dev)).cpu().numpy()
+    grid(fake, os.path.join(a.out, "grid_final.png"))
+    torch.save({"g":g.state_dict(),"d":d.state_dict(),"step":a.steps,"zdim":a.zdim}, os.path.join(a.out,"final.pt"))
+    lg(f"DONE final step {a.steps}")
+    lg("files: " + ", ".join(sorted(os.listdir(a.out))))
+    logf.close()
+
+if __name__=="__main__":
+    main()