Upload Splat_trainer2.py

Splat_trainer2.py

@@ -0,0 +1,441 @@
+#!/usr/bin/env python3
+# splat_trainer2.py — the fast trainer (faces folder -> better splat_decoder.onnx)
+#
+# Same architecture as splat_generator.py (latent 128, Gabor packets, anchor
+# grid, complex phase head) so every existing tool — splat_cv5, probe, surf,
+# atlas, zoom — works on the new model unchanged. What changed is SPEED:
+#
+#   1. CACHE ONCE. The old trainer decoded 200k JPEGs every epoch — that was
+#      the real bottleneck, not the GPU. First run builds faces_cache_S.npy
+#      (uint8, center-cropped, resized) with threaded cv2. Every later run
+#      starts in seconds.
+#   2. DATASET LIVES ON THE GPU. 202k x 96x96x3 uint8 = 5.6 GB -> fits a 12GB
+#      card next to the model (64px = 2.5 GB). Batches are fancy-indexed on
+#      device; there is NO DataLoader, no workers, no H2D copy per step.
+#      Falls back to pinned CPU memory automatically if it doesn't fit.
+#   3. VECTORIZED RENDERER. The per-channel python loop is now shared-carrier
+#      multiply-sums per chunk (env*cos and env*sin are computed once, not three times).
+#      Verified equal to the old loop renderer to float tolerance in --smoke.
+#   4. STEPS, NOT EPOCHS. VAEs converge per gradient step; random batches
+#      from the resident tensor, cosine LR with warmup, KL beta ramped in
+#      steps. --steps 30000 at batch 96 sees ~2.9M images (14 "epochs") in
+#      roughly the wall time the old loop needed for 2.
+#   5. bf16 autocast for encoder/decoder (renderer stays fp32, as always),
+#      fused Adam when available, gradient checkpointing OFF by default
+#      (it halves VRAM but doubles renderer compute — flag it back on only
+#      if you OOM).
+#
+#   python splat_trainer2.py --data_dir E:/path/to/faces          # train
+#   python splat_trainer2.py --export                             # -> onnx
+#   python splat_trainer2.py --smoke                              # CPU test
+#
+# The export writes splat_decoder.onnx with the exact input/output names
+# ("z_latent" / "rendered_image", opset 17, dynamic batch) the cv5 tools use.
+#
+# HONESTY: --smoke was run end-to-end (train -> export -> cv.dnn reload ->
+# torch/ONNX parity) on CPU in the sandbox. The full-speed GPU path (bf16,
+# fused Adam, resident-tensor indexing) follows the same code but its
+# throughput numbers are yours to measure. PerceptionLab discipline: do not
+# hype, do not lie, just show.
+
+import argparse, glob, math, os, sys, time
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+K = 11               # dpx,dpy,ls,th,lf + (a,b) x 3 channels
+LATENT = 128         # fixed: every downstream tool assumes it
+
+# ======================================================================
+# 1) preprocessing cache: faces folder -> uint8 npy, once
+# ======================================================================
+def build_cache(data_dir, size, cache_path):
+    import cv2 as cv
+    from concurrent.futures import ThreadPoolExecutor
+    exts = ("*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp")
+    paths = sorted(p for e in exts for p in glob.glob(os.path.join(data_dir, e)))
+    if not paths:
+        raise RuntimeError(f"no images in {data_dir}")
+    n = len(paths)
+    print(f"caching {n} images at {size}px -> {cache_path} (one time)")
+    arr = np.lib.format.open_memmap(cache_path, mode="w+", dtype=np.uint8,
+                                    shape=(n, size, size, 3))
+    def work(i):
+        im = cv.imread(paths[i], cv.IMREAD_COLOR)
+        if im is None:
+            return i, False
+        h, w = im.shape[:2]
+        s = min(h, w)
+        im = im[(h - s) // 2:(h + s) // 2, (w - s) // 2:(w + s) // 2]
+        im = cv.resize(im, (size, size), interpolation=cv.INTER_AREA)
+        arr[i] = im[:, :, ::-1]                        # BGR -> RGB
+        return i, True
+    t0, done = time.time(), 0
+    with ThreadPoolExecutor(max_workers=os.cpu_count()) as ex:
+        for i, ok in ex.map(work, range(n)):
+            done += 1
+            if done % 20000 == 0:
+                r = done / (time.time() - t0)
+                print(f"  {done}/{n}  ({r:.0f} img/s, eta {(n-done)/r/60:.1f} min)")
+    arr.flush()
+    print(f"cache built in {(time.time()-t0)/60:.1f} min")
+
+def load_resident(cache_path, dev):
+    """Whole dataset as a uint8 tensor, on GPU if it fits."""
+    a = np.load(cache_path, mmap_mode="r")
+    t = torch.from_numpy(np.ascontiguousarray(a))
+    if dev.type == "cuda":
+        need = t.numel()
+        free, _ = torch.cuda.mem_get_info()
+        if need < free - 3e9:                          # leave 3GB for training
+            t = t.to(dev)
+            print(f"dataset resident on GPU: {need/1e9:.2f} GB, {len(t)} images")
+            return t
+        t = t.pin_memory()
+        print(f"dataset pinned on CPU ({need/1e9:.2f} GB too big for VRAM)")
+    return t
+
+def batch_from(data, idx, dev):
+    x = data[idx]
+    if x.device != dev:
+        x = x.to(dev, non_blocking=True)
+    return x.permute(0, 3, 1, 2).float().div_(255.0)
+
+# ======================================================================
+# 2) model — identical math to splat_generator.py, faster renderer
+# ======================================================================
+class GaborRenderer(nn.Module):
+    def __init__(self, image_size=96, num_packets=256, chunk=64, use_checkpoint=False):
+        super().__init__()
+        self.H = self.W = image_size
+        self.N, self.chunk, self.use_checkpoint = num_packets, chunk, use_checkpoint
+        gy, gx = torch.meshgrid(torch.linspace(0, 1, image_size),
+                                torch.linspace(0, 1, image_size), indexing="ij")
+        self.register_buffer("GX", gx[None, None].contiguous())
+        self.register_buffer("GY", gy[None, None].contiguous())
+        side = int(math.ceil(math.sqrt(num_packets)))
+        ax = torch.linspace(0.08, 0.92, side)
+        anch = torch.stack(torch.meshgrid(ax, ax, indexing="ij"), -1).reshape(-1, 2)[:num_packets]
+        anch = torch.clamp(anch, 1e-3, 1 - 1e-3)
+        self.register_buffer("anchor_logit", torch.log(anch / (1 - anch)))
+
+    def activate(self, raw):
+        px = torch.sigmoid(self.anchor_logit[:, 0][None] + raw[..., 0])
+        py = torch.sigmoid(self.anchor_logit[:, 1][None] + raw[..., 1])
+        sigma = 0.012 + 0.14 * torch.sigmoid(raw[..., 2])
+        theta = raw[..., 3]
+        freq = 1.0 + 15.0 * torch.sigmoid(raw[..., 4])
+        coeff = torch.tanh(raw[..., 5:11]).reshape(*raw.shape[:2], 3, 2)
+        return px, py, sigma, theta, freq, coeff
+
+    def _chunk(self, px, py, sigma, theta, freq, coeff):
+        """Vectorized: env*cos / env*sin once, channels via one einsum each."""
+        px_ = px[..., None, None]; py_ = py[..., None, None]
+        s_ = sigma[..., None, None]; th = theta[..., None, None]
+        f_ = freq[..., None, None]
+        dx = self.GX - px_; dy = self.GY - py_
+        xr = dx * torch.cos(th) + dy * torch.sin(th)
+        env = torch.exp(-(dx * dx + dy * dy) / (2 * s_ * s_))
+        ec = env * torch.cos(2 * math.pi * f_ * xr)          # (B,n,H,W)
+        es = env * torch.sin(2 * math.pi * f_ * xr)
+        a, b = coeff[..., 0], coeff[..., 1]                  # (B,n,3)
+        # per-channel multiply-sum: ec/es are still computed ONCE (the speed
+        # win over the old loop), and the graph is pure Mul+ReduceSum+Stack —
+        # no Einsum, no dynamic Reshape — so it runs bit-identically on cv2
+        # 4.x legacy dnn AND cv5 ENGINE_NEW, at any batch size
+        chans = [(a[:, :, c, None, None] * ec).sum(1)
+                 - (b[:, :, c, None, None] * es).sum(1) for c in range(3)]
+        return torch.stack(chans, dim=1)
+
+    def forward(self, raw):
+        raw = raw.float()                                    # fp32 always
+        px, py, sigma, theta, freq, coeff = self.activate(raw)
+        out = None                       # no zeros(batch,...): keeps the ONNX
+        for i in range(0, self.N, self.chunk):   # graph free of ConstantOfShape
+            sl = slice(i, i + self.chunk)
+            args = (px[:, sl], py[:, sl], sigma[:, sl],
+                    theta[:, sl], freq[:, sl], coeff[:, sl])
+            if self.use_checkpoint and self.training:
+                from torch.utils.checkpoint import checkpoint
+                c = checkpoint(self._chunk, *args, use_reentrant=False)
+            else:
+                c = self._chunk(*args)
+            out = c if out is None else out + c
+        return torch.sigmoid(out)
+
+class Encoder(nn.Module):
+    def __init__(self, image_size=96, latent=LATENT, ch=32):
+        super().__init__()
+        layers, c_in, sz, c = [], 3, image_size, ch
+        while sz > 4:
+            layers += [nn.Conv2d(c_in, c, 4, 2, 1), nn.BatchNorm2d(c),
+                       nn.LeakyReLU(0.2, True)]
+            c_in, sz, c = c, sz // 2, min(c * 2, 512)
+        self.conv = nn.Sequential(*layers)
+        self.flat = c_in * sz * sz
+        self.fc_mu = nn.Linear(self.flat, latent)
+        self.fc_lv = nn.Linear(self.flat, latent)
+    def forward(self, x):
+        h = self.conv(x).flatten(1)
+        return self.fc_mu(h), self.fc_lv(h)
+
+class Decoder(nn.Module):
+    def __init__(self, latent=LATENT, num_packets=256, hidden=512):
+        super().__init__()
+        self.N = num_packets
+        self.net = nn.Sequential(
+            nn.Linear(latent, hidden), nn.LeakyReLU(0.2, True),
+            nn.Linear(hidden, hidden), nn.LeakyReLU(0.2, True),
+            nn.Linear(hidden, num_packets * K))
+        nn.init.zeros_(self.net[-1].bias)
+        self.net[-1].weight.data *= 0.1
+    def forward(self, z):
+        return self.net(z).view(-1, self.N, K)
+
+class SplatVAE(nn.Module):
+    def __init__(self, image_size=96, num_packets=256, chunk=64, ckpt=False):
+        super().__init__()
+        self.enc = Encoder(image_size)
+        self.dec = Decoder(LATENT, num_packets)
+        self.ren = GaborRenderer(image_size, num_packets, chunk, ckpt)
+        self.latent = LATENT
+
+def kl(mu, lv):
+    return -0.5 * torch.mean(torch.sum(1 + lv - mu.pow(2) - lv.exp(), dim=1))
+
+# ======================================================================
+# 3) training — steps, resident data, bf16, cosine LR
+# ======================================================================
+def train(args, dev):
+    cache = os.path.join(args.out, f"faces_cache_{args.image_size}.npy")
+    os.makedirs(args.out, exist_ok=True)
+    if not os.path.exists(cache):
+        build_cache(args.data_dir, args.image_size, cache)
+    data = load_resident(cache, dev)
+    n = len(data)
+
+    model = SplatVAE(args.image_size, args.num_packets, args.chunk,
+                     args.checkpointing).to(dev)
+    if args.resume and os.path.exists(args.resume):
+        model.load_state_dict(torch.load(args.resume, map_location=dev)["sd"])
+        print("resumed", args.resume)
+    print(f"params {sum(p.numel() for p in model.parameters())/1e6:.2f}M  "
+          f"steps {args.steps}  batch {args.batch}  res {args.image_size}")
+
+    fused = dev.type == "cuda"
+    opt = torch.optim.Adam(model.parameters(), lr=args.lr, fused=fused)
+    warm = max(1, args.steps // 50)
+    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(
+        (s + 1) / warm, 0.5 * (1 + math.cos(math.pi * s / args.steps))))
+    use_bf16 = dev.type == "cuda" and torch.cuda.is_bf16_supported()
+    print(f"autocast bf16: {use_bf16}  fused adam: {fused}  "
+          f"checkpointing: {args.checkpointing}")
+
+    g = torch.Generator(device="cpu").manual_seed(0)
+    fixed_idx = torch.randint(0, n, (32,), generator=g)
+    z_fixed = torch.randn(64, LATENT, device=dev)
+    logf = open(os.path.join(args.out, "loss.csv"), "a")
+    t0, run_rec, run_kl, last = time.time(), 0.0, 0.0, 0
+    model.train()
+    for step in range(1, args.steps + 1):
+        idx = torch.randint(0, n, (args.batch,), generator=g)
+        x = batch_from(data, idx, dev)
+        beta = args.beta * min(1.0, step / max(1, args.beta_warmup_steps))
+        opt.zero_grad(set_to_none=True)
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_bf16):
+            mu, lv = model.enc(x)
+            z = mu + torch.randn_like(mu) * torch.exp(0.5 * lv)
+            raw = model.dec(z)
+        recon = model.ren(raw)                          # fp32 renderer
+        rec = F.mse_loss(recon, x)
+        # floater penalty: charge amplitude carried by needle-thin envelopes.
+        # the floater strategy = sigma -> min, amp -> max (a bright orphan dot
+        # that patches one pixel). amp^2 * max(SIGMA_REF/sigma - 1, 0) prices
+        # point-brightness: zero cost above SIGMA_REF, growing cost as the
+        # envelope collapses toward the floor. gamma_floater=0 disables.
+        if args.gamma_floater > 0:
+            _, _, sg, _, _, cf = model.ren.activate(raw.float())
+            amp2 = cf.pow(2).sum(dim=(-1, -2))          # (B,N) per-packet energy
+            flo = (amp2 * (args.sigma_ref / sg - 1.0).clamp(min=0)).mean()
+        else:
+            flo = torch.zeros((), device=x.device)
+        loss = rec + beta * kl(mu, lv) + args.gamma_floater * flo
+        loss.backward()
+        nn.utils.clip_grad_norm_(model.parameters(), 5.0)
+        opt.step(); sched.step()
+        run_rec += rec.item(); run_kl += kl(mu, lv).item()
+
+        if step % args.log_every == 0 or step == args.steps:
+            nb = step - last; last = step
+            ips = nb * args.batch / (time.time() - t0); t0 = time.time()
+            psnr = 10 * math.log10(1.0 / max(run_rec / nb, 1e-9))
+            print(f"step {step:6d}/{args.steps}  rec {run_rec/nb:.4f} "
+                  f"(PSNR {psnr:4.1f})  kl {run_kl/nb:7.1f}  beta {beta:.2f}  "
+                  f"lr {sched.get_last_lr()[0]:.2e}  {ips:6.0f} img/s")
+            logf.write(f"{step},{run_rec/nb:.6f},{run_kl/nb:.6f}\n"); logf.flush()
+            run_rec = run_kl = 0.0
+            model.eval()
+            with torch.no_grad():
+                torch.save({"sd": model.state_dict(),
+                            "image_size": args.image_size,
+                            "num_packets": args.num_packets},
+                           os.path.join(args.out, "model2.pt"))
+                fx = batch_from(data, fixed_idx, dev)
+                mu, _ = model.enc(fx)
+                rc = model.ren(model.dec(mu))
+                grid(torch.cat([fx, rc], 0),
+                     os.path.join(args.out, f"recon_{step:06d}.png"))
+                grid(model.ren(model.dec(z_fixed)),
+                     os.path.join(args.out, f"sample_{step:06d}.png"))
+            model.train()
+    print("done ->", os.path.join(args.out, "model2.pt"),
+          " | now: python splat_trainer2.py --export")
+
+def grid(t, path, nrow=8):
+    import cv2 as cv
+    t = t.clamp(0, 1).cpu().numpy()
+    n, _, h, w = t.shape
+    rows = int(math.ceil(n / nrow))
+    g = np.zeros((rows * h, nrow * w, 3), np.float32)
+    for i in range(n):
+        r, c = divmod(i, nrow)
+        g[r*h:(r+1)*h, c*w:(c+1)*w] = np.transpose(t[i], (1, 2, 0))
+    cv.imwrite(path, (g[:, :, ::-1] * 255).astype(np.uint8))
+
+# ======================================================================
+# 4) ONNX export — same contract as the cv5 tools expect
+# ======================================================================
+class ExportHead(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.dec, self.ren = model.dec, model.ren
+        self.ren.use_checkpoint = False
+    def forward(self, z):
+        return self.ren(self.dec(z))
+
+def export(args, dev):
+    ck = torch.load(os.path.join(args.out, "model2.pt"), map_location="cpu")
+    model = SplatVAE(ck["image_size"], ck["num_packets"], args.chunk)
+    model.load_state_dict(ck["sd"]); model.eval()
+    head = ExportHead(model)
+    dummy = torch.randn(1, LATENT)
+    out = args.onnx or "splat_decoder.onnx"
+    torch.onnx.export(head, dummy, out, export_params=True, opset_version=17,
+                      do_constant_folding=True, input_names=["z_latent"],
+                      output_names=["rendered_image"],
+                      dynamic_axes={"z_latent": {0: "batch"},
+                                    "rendered_image": {0: "batch"}},
+                      dynamo=False)
+    mb = os.path.getsize(out) / 1e6
+    print(f"exported {out} ({mb:.1f} MB, {ck['image_size']}px, "
+          f"{ck['num_packets']} packets) — drop-in for the cv5 tools")
+
+# ======================================================================
+# 5) smoke — CPU end-to-end: loop-vs-einsum parity, train, export, cv.dnn parity
+# ======================================================================
+def smoke():
+    ok = True
+    def check(name, cond, note=""):
+        nonlocal ok; ok &= bool(cond)
+        print(f"  [{'PASS' if cond else 'FAIL'}] {name} {note}")
+    torch.manual_seed(0)
+    dev = torch.device("cpu")
+
+    # (a) vectorized renderer == original per-channel loop renderer
+    ren = GaborRenderer(32, 16, chunk=8)
+    raw = torch.randn(2, 16, K) * 0.5
+    with torch.no_grad():
+        fast = ren(raw)
+        px, py, sg, th, fq, cf = ren.activate(raw.float())
+        outs = []
+        for c in range(3):                              # the old loop, verbatim
+            px_ = px[..., None, None]; py_ = py[..., None, None]
+            s_ = sg[..., None, None]; t_ = th[..., None, None]
+            f_ = fq[..., None, None]
+            dx = ren.GX - px_; dy = ren.GY - py_
+            xr = dx * torch.cos(t_) + dy * torch.sin(t_)
+            env = torch.exp(-(dx*dx + dy*dy) / (2*s_*s_))
+            a = cf[:, :, c, 0][..., None, None]; b = cf[:, :, c, 1][..., None, None]
+            outs.append((env * (a*torch.cos(2*math.pi*f_*xr)
+                              - b*torch.sin(2*math.pi*f_*xr))).sum(1))
+        slow = torch.sigmoid(torch.stack(outs, 1))
+    err = (fast - slow).abs().max().item()
+    check("einsum renderer == loop renderer", err < 1e-5, f"max|d| {err:.2e}")
+
+    # (b) tiny synthetic cache + short training run: loss must fall
+    import tempfile, cv2 as cv
+    tmp = tempfile.mkdtemp()
+    imdir = os.path.join(tmp, "imgs"); os.makedirs(imdir)
+    rng = np.random.default_rng(0)
+    for i in range(24):
+        im = np.zeros((40, 36, 3), np.uint8)
+        cv.circle(im, (rng.integers(8, 28), rng.integers(8, 32)),
+                  rng.integers(4, 10), tuple(int(v) for v in rng.integers(60, 255, 3)), -1)
+        cv.imwrite(os.path.join(imdir, f"{i:03d}.png"), im)
+    a = argparse.Namespace(
+        data_dir=imdir, out=tmp, image_size=32, num_packets=16, chunk=8,
+        batch=8, steps=60, lr=3e-3, beta=1e-4, beta_warmup_steps=30,
+        log_every=30, resume="", checkpointing=False, gamma_floater=0.02,
+        sigma_ref=0.03, onnx=os.path.join(tmp, "t.onnx"))
+    import io, contextlib
+    buf = io.StringIO()
+    with contextlib.redirect_stdout(buf):
+        train(a, dev)
+    lines = [l for l in buf.getvalue().splitlines() if l.startswith("step")]
+    r0 = float(lines[0].split("rec")[1].split("(")[0])
+    r1 = float(lines[-1].split("rec")[1].split("(")[0])
+    check("training loss falls", r1 < r0, f"{r0:.4f} -> {r1:.4f}")
+    check("cache built", os.path.exists(os.path.join(tmp, "faces_cache_32.npy")))
+
+    # (c) export + cv.dnn reload + parity with torch
+    with contextlib.redirect_stdout(buf):
+        export(a, dev)
+    check("onnx written", os.path.exists(a.onnx))
+    ck = torch.load(os.path.join(tmp, "model2.pt"), map_location="cpu")
+    m = SplatVAE(32, 16, 8); m.load_state_dict(ck["sd"]); m.eval()
+    z = torch.randn(3, LATENT)
+    with torch.no_grad():
+        want = ExportHead(m)(z).numpy()
+    net = cv.dnn.readNetFromONNX(a.onnx)
+    net.setInput(z.numpy(), "z_latent")
+    got = net.forward("rendered_image")
+    err = float(np.abs(got - want).max())
+    check("cv.dnn output == torch output", err < 1e-4,
+          f"max|d| {err:.2e}, batch of 3 through dynamic axis")
+    print("smoke:", "ALL PASS" if ok else "FAILURES ABOVE")
+    return 0 if ok else 1
+
+# ======================================================================
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data_dir", default="./faces")
+    ap.add_argument("--out", default="./runs/splat2")
+    ap.add_argument("--image_size", type=int, default=96)
+    ap.add_argument("--num_packets", type=int, default=256)
+    ap.add_argument("--chunk", type=int, default=64)
+    ap.add_argument("--batch", type=int, default=96)
+    ap.add_argument("--steps", type=int, default=30000)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--beta", type=float, default=1.0)
+    ap.add_argument("--beta_warmup_steps", type=int, default=3000)
+    ap.add_argument("--gamma_floater", type=float, default=0.02,
+                    help="anti-floater energy penalty (0 = off)")
+    ap.add_argument("--sigma_ref", type=float, default=0.03,
+                    help="envelopes thinner than this pay the penalty")
+    ap.add_argument("--log_every", type=int, default=250)
+    ap.add_argument("--resume", default="")
+    ap.add_argument("--checkpointing", action="store_true",
+                    help="halve VRAM, double renderer compute (only if OOM)")
+    ap.add_argument("--export", action="store_true")
+    ap.add_argument("--onnx", default=None)
+    ap.add_argument("--smoke", action="store_true")
+    args = ap.parse_args()
+    if args.smoke:
+        sys.exit(smoke())
+    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("device:", dev)
+    if args.export:
+        export(args, dev)
+    else:
+        train(args, dev)