File size: 4,046 Bytes
35ec68e | 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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | """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 # identity folders used to train the proxy
PROXY_ITERS = 30000 # paper: 30k iterations
BATCH = 128
COV_SAMPLES = 20000 # samples used to estimate the gradient covariance
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)
# ---- gradient covariance over the channel-averaged reconstruction grad ----
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] # (B,C,H,W)
gbar.append(g.mean(dim=(2, 3)).detach().cpu()) # (B,C)
seen += x.shape[0]
if seen >= COV_SAMPLES:
break
gbar = torch.cat(gbar, 0) # (N,C)
mu = gbar.mean(0, keepdim=True)
gc = gbar - mu
sigma = (gc.T @ gc) / (gbar.shape[0] - 1) # (C,C)
evals, evecs = torch.linalg.eigh(sigma) # ascending
order = torch.argsort(evals, descending=True)
evals, evecs = evals[order], evecs[:, order]
U_rec = evecs[:, :CFG["r"]] # (C,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()
|