| 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) |
| imgs = (imgs - 0.5)/0.5 |
| 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) |
| |
| 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() |
| |
| 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() |
| |
| 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() |
|
|
