| """RSSD: train the proxy inversion network Phi, then estimate the |
| reconstruction-sensitive subspace from the gradient covariance (paper Sec. IV-B). |
| |
| Outputs checkpoints/rssd.pt with P_rec, P_N (C x C), U_rec, eigenvalues, and the |
| trained proxy checkpoint checkpoints/proxy_unet.pt. |
| """ |
| import argparse |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data import DataLoader |
| import wandb |
|
|
| from models import load_ir50, ClientEncoder, UNetInverter |
| from data import MS1MDataset |
| from common import ADAFACE_CKPT, CKPT, CFG, DEVICE |
|
|
| PROXY_IDS = 5000 |
| PROXY_ITERS = 30000 |
| BATCH = 128 |
| COV_SAMPLES = 20000 |
|
|
|
|
| def infinite(loader): |
| while True: |
| for batch in loader: |
| yield batch |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ms1m_root", required=True) |
| args = ap.parse_args() |
|
|
| CKPT.mkdir(exist_ok=True) |
| dev = DEVICE |
| wandb.init(project="decoyface-repro", name="rssd", mode="online", config=CFG) |
|
|
| ir50 = load_ir50(ADAFACE_CKPT, dev) |
| Fc = ClientEncoder(ir50).to(dev).eval() |
| for p in Fc.parameters(): |
| p.requires_grad_(False) |
|
|
| ds = MS1MDataset(args.ms1m_root, PROXY_IDS) |
| loader = DataLoader(ds, batch_size=BATCH, shuffle=True, num_workers=8, |
| drop_last=True, pin_memory=True, persistent_workers=True) |
| it = infinite(loader) |
|
|
| proxy = UNetInverter().to(dev) |
| opt = torch.optim.Adam(proxy.parameters(), lr=1e-3) |
|
|
| proxy_path = CKPT / "proxy_unet.pt" |
| resume_path = CKPT / "proxy_unet_resume.pt" |
| start = 0 |
| if resume_path.exists(): |
| ck = torch.load(resume_path, map_location=dev) |
| proxy.load_state_dict(ck["model"]) |
| opt.load_state_dict(ck["opt"]) |
| start = ck["step"] |
| print(f"resuming proxy from step {start}", flush=True) |
| for step in range(start, PROXY_ITERS): |
| x, _ = next(it) |
| x = x.to(dev, non_blocking=True) |
| with torch.no_grad(): |
| z = Fc(x) |
| rec = proxy(z) |
| loss = F.l1_loss(rec, x) |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| opt.step() |
| if step % 50 == 0: |
| wandb.log({"proxy/l1": loss.item()}, step=step) |
| if step % 2000 == 0: |
| print(f"proxy step {step} l1 {loss.item():.4f}", flush=True) |
| torch.save({"step": step, "model": proxy.state_dict(), |
| "opt": opt.state_dict()}, resume_path) |
| torch.save(proxy.state_dict(), proxy_path) |
|
|
| |
| proxy.eval() |
| gbar = [] |
| seen = 0 |
| for x, _ in DataLoader(ds, batch_size=64, shuffle=True, num_workers=8): |
| x = x.to(dev) |
| z = Fc(x).detach().clone().requires_grad_(True) |
| loss = F.l1_loss(proxy(z), x) |
| g = torch.autograd.grad(loss, z)[0] |
| gbar.append(g.mean(dim=(2, 3)).detach().cpu()) |
| seen += x.shape[0] |
| if seen >= COV_SAMPLES: |
| break |
| gbar = torch.cat(gbar, 0) |
| mu = gbar.mean(0, keepdim=True) |
| gc = gbar - mu |
| sigma = (gc.T @ gc) / (gbar.shape[0] - 1) |
| evals, evecs = torch.linalg.eigh(sigma) |
| order = torch.argsort(evals, descending=True) |
| evals, evecs = evals[order], evecs[:, order] |
| U_rec = evecs[:, :CFG["r"]] |
| P_rec = U_rec @ U_rec.T |
| P_N = torch.eye(CFG["C"]) - P_rec |
|
|
| torch.save({"P_rec": P_rec, "P_N": P_N, "U_rec": U_rec, |
| "evals": evals, "sigma": sigma}, CKPT / "rssd.pt") |
| ev = evals / evals.sum() |
| print("top-8 eigenvalue energy:", ev[:8].tolist(), flush=True) |
| print(f"RSS energy captured by top-{CFG['r']}: " |
| f"{ev[:CFG['r']].sum().item():.4f}", flush=True) |
| wandb.log({"rssd/rss_energy": ev[:CFG["r"]].sum().item()}) |
| wandb.finish() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
