File size: 4,598 Bytes
3ba6625 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | 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()
|