| |
| |
| |
| |
| |
| import os, sys, time, math, json |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data import Dataset, DataLoader |
| from PIL import Image |
| sys.path.insert(0, "/home/z/my-project/scripts/os4") |
| from os4_models import FrameGenerator |
|
|
| DATA = "/home/z/my-project/data/einkos4" |
| CKPT = "/home/z/my-project/checkpoints4/gen_ckpt.pt" |
| N_TOTAL = 3200 |
| PHASE0_STEPS = 6000 |
| PHASE1_STEPS = 900 |
|
|
| torch.set_num_threads(2) |
|
|
|
|
| class Os4DS(Dataset): |
| def __init__(self, phase, idxs): |
| self.phase = phase |
| self.idxs = idxs |
| self.fm = np.memmap(os.path.join(DATA, "frames.mem"), dtype=np.uint8, mode="r", |
| shape=(N_TOTAL, 832, 384, 3)) |
| self.states = np.load(os.path.join(DATA, "states.npy")) |
|
|
| def __len__(self): |
| return len(self.idxs) |
|
|
| def __getitem__(self, i): |
| idx = self.idxs[i] |
| fr = np.asarray(self.fm[idx]) |
| hint = np.asarray(Image.open(os.path.join(DATA, "hint", "%06d.png" % idx)).convert("RGB")) |
| if self.phase == 0: |
| fr = fr.reshape(416, 2, 192, 2, 3).mean(axis=(1, 3)) |
| hint = hint.reshape(208, 2, 96, 2, 3).mean(axis=(1, 3)) |
| t = torch.from_numpy(np.ascontiguousarray(fr)).permute(2, 0, 1).float() / 255.0 |
| h = torch.from_numpy(np.ascontiguousarray(hint)).permute(2, 0, 1).float() / 255.0 |
| s = torch.from_numpy(self.states[idx].copy()) |
| return h, s, t |
|
|
|
|
| def sobel_mag(x): |
| kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3) |
| ky = kx.transpose(2, 3).contiguous() |
| g = x.mean(1, keepdim=True) |
| gx = F.conv2d(g, kx, padding=1) |
| gy = F.conv2d(g, ky, padding=1) |
| return torch.sqrt(gx * gx + gy * gy + 1e-6) |
|
|
|
|
| def loss_fn(pred, tgt): |
| l1 = (pred - tgt).abs().mean() |
| e = (sobel_mag(pred) - sobel_mag(tgt)).abs().mean() |
| return l1 + 0.25 * e, l1.item(), e.item() |
|
|
|
|
| def make_val(idxs_all): |
| return [i for i in idxs_all if i % 16 == 5][:24] |
|
|
|
|
| def main(): |
| dev = "cpu" |
| model = FrameGenerator().to(dev) |
| opt = None |
| step = 0 |
| phase = 0 |
| best = 1e9 |
| if os.path.exists(CKPT): |
| ck = torch.load(CKPT, map_location=dev, weights_only=False) |
| model.load_state_dict(ck["m"]) |
| phase = ck["phase"]; step = ck["step"]; best = ck.get("best", 1e9) |
| print("resume step %d phase %d best %.4f" % (step, phase, best), flush=True) |
| lr = 2.5e-4 if phase == 0 else 8e-5 |
| lr_end = 1.2e-4 if phase == 0 else 2.5e-5 |
| extra = int(sys.argv[1]) if len(sys.argv) > 1 else 0 |
| p0 = step if phase == 0 else 0 |
| p1 = step - PHASE0_STEPS if phase == 1 else 0 |
| bs = 6 if phase == 0 else 2 |
| total = PHASE0_STEPS if phase == 0 else PHASE1_STEPS + extra |
| done = p0 if phase == 0 else p1 |
|
|
| idxs = list(range(N_TOTAL)) |
| import random |
| random.Random(7).shuffle(idxs) |
| ds = Os4DS(phase, idxs) |
| dl = DataLoader(ds, batch_size=bs, shuffle=True, num_workers=2, |
| drop_last=True, persistent_workers=False, prefetch_factor=2) |
| opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5) |
| if os.path.exists(CKPT) and phase == ck["phase"]: |
| try: |
| opt.load_state_dict(ck["o"]) |
| except Exception: |
| pass |
|
|
| t0 = time.time() |
| ep = 0 |
| while True: |
| for hb, sb, tb in dl: |
| |
| prog = min(1.0, done / float(total)) |
| cur = lr_end + 0.5 * (lr - lr_end) * (1 + math.cos(math.pi * prog)) |
| for g in opt.param_groups: |
| g["lr"] = cur |
| pred = model(hb, sb) |
| loss, l1v, ev = loss_fn(pred, tb) |
| opt.zero_grad() |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| opt.step() |
| done += 1; step += 1 |
| if done % 30 == 0: |
| tmp = CKPT + ".tmp" |
| torch.save({"m": model.state_dict(), "o": opt.state_dict(), |
| "step": step, "phase": phase, "best": best}, tmp) |
| os.replace(tmp, CKPT) |
| if done % 150 == 0: |
| el = time.time() - t0 |
| print("p%d step %d/%d loss %.4f (l1 %.4f e %.4f) lr %.1e %.2f s/it" % |
| (phase, done, total, loss.item(), l1v, ev, cur, el / 150.0), flush=True) |
| t0 = time.time() |
| if done % 600 == 0 or done == total: |
| model.eval() |
| ps = [] |
| with torch.no_grad(): |
| for vi in make_val(idxs): |
| h, s, t = ds[vi] |
| p = model(h.unsqueeze(0), s.unsqueeze(0)) |
| mse = ((p - t.unsqueeze(0)) ** 2).mean().item() |
| ps.append(10 * math.log10(1.0 / max(1e-8, mse))) |
| v = float(np.mean(ps)) |
| print(" [val] phase %d step %d PSNR %.2f dB" % (phase, done, v), flush=True) |
| if v < best: |
| best = v |
| torch.save({"m": model.state_dict(), "step": step, "phase": phase, |
| "best": best, "opt": None}, CKPT.replace(".pt", "_best.pt")) |
| model.train() |
| if done >= total: |
| if phase == 0: |
| phase = 1 |
| ck = torch.load(CKPT, map_location=dev, weights_only=False) |
| ck["phase"] = 1; ck["step"] = PHASE0_STEPS |
| tmp = CKPT + ".tmp" |
| torch.save(ck, tmp); os.replace(tmp, CKPT) |
| print("PHASE 0 COMPLETE -> phase 1 (app-res fine-tune)", flush=True) |
| return |
| else: |
| torch.save({"m": model.state_dict(), "step": step, "phase": 2, |
| "best": best}, CKPT.replace(".pt", "_final.pt")) |
| print("TRAINING DONE step %d best-val %.2f dB" % (step, best), flush=True) |
| return |
| ep += 1 |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
