Clean DecoyFace implementation and checkpoints

README.md

@@ -0,0 +1,70 @@
+---
+library_name: pytorch
+tags:
+- face-recognition
+- privacy
+---
+
+# DecoyFace
+
+Implementation of DecoyFace split face recognition privacy protection:
+client-side Decoy-Guided Coherent Mixing (DGCM) and the server-side Authorized Canonicalization Module (ACM),
+together with the scripts that build the Reconstruction-Sensitive Subspace Decomposition (RSSD) statistics and train the server backbone.
+
+The client encoder `F_c` is AdaFace IR-50 (MS1MV2) stem + stage 1 (output 64x56x56);
+the server backbone `F_s` is the rest of the IR-50.
+
+## Files
+
+| File | Content |
+| --- | --- |
+| `decoyface.py` | `DecoyProtector`: `protect(z_A, z_B)` (client, DGCM) and `server_recover(z_t)` (server, ACM) |
+| `models.py` | IR-50, its split into `ClientEncoder` (F_c) / `ServerBackbone` (F_s), and the proxy inversion U-Net used by RSSD |
+| `common.py` | DecoyFace hyper-parameters `CFG`, fixed channel permutation, checkpoint paths |
+| `data.py` | MS1M identity-folder dataset |
+| `build_rssd.py` | Trains the proxy inversion U-Net, estimates the reconstruction-sensitive subspace -> `checkpoints/rssd.pt` |
+| `build_bank.py` | Builds the auxiliary complementary bank and `mu_rec` -> `checkpoints/bank.pt` |
+| `train_defender.py` | Trains F_s + ArcFace head on canonicalized protected features -> `checkpoints/defender.pt` |
+
+## Checkpoints
+
+| File | Content |
+| --- | --- |
+| `checkpoints/adaface_ir50_ms1mv2.backbone.pth` | AdaFace IR-50 MS1MV2 backbone state dict (from https://github.com/mk-minchul/AdaFace) |
+| `checkpoints/rssd.pt` | `P_rec`, `P_N` (64x64), `U_rec`, `evals`, `sigma` |
+| `checkpoints/bank.pt` | `aux_bank` (5000x64x56x56, fp16), `mu_rec` (64x56x56), `perm` (64) |
+
+`proxy_unet.pt` and `defender.pt` are not included.
+
+## Usage
+
+```python
+import torch
+from common import ADAFACE_CKPT, CKPT, CFG
+from models import load_ir50, ClientEncoder
+from decoyface import DecoyProtector
+
+Fc = ClientEncoder(load_ir50(ADAFACE_CKPT, "cuda")).cuda().eval()
+rssd = torch.load(CKPT / "rssd.pt", map_location="cuda")
+bank = torch.load(CKPT / "bank.pt", map_location="cuda")
+protector = DecoyProtector(rssd["P_rec"], rssd["P_N"], bank["mu_rec"],
+                           bank["aux_bank"].float(), bank["perm"], CFG, "cuda")
+
+# x, decoy: (B,3,112,112) RGB in [0,1], ArcFace-aligned; decoy is a different identity
+with torch.no_grad():
+    z_t = protector.protect(Fc(x), Fc(decoy))   # transmitted by the client
+    z_can = protector.server_recover(z_t)        # fed to F_s on the server
+```
+
+## Rebuilding
+
+MS1M root: identity folders named by integer, each holding 112x112 aligned `.jpg` faces.
+Identities used: proxy U-Net first 5000, auxiliary bank 50000-60000, defender first 30000 (`--num_ids`).
+
+```bash
+python build_rssd.py --ms1m_root /path/to/ms1m
+python build_bank.py --ms1m_root /path/to/ms1m
+python train_defender.py --ms1m_root /path/to/ms1m
+```
+
+Scripts log to Weights & Biases (project `decoyface-repro`).

build_bank.py

@@ -0,0 +1,63 @@
+"""Build the auxiliary complementary bank D_N = {P_N F_c(y_j)} and the fixed
+reconstruction-sensitive anchor mu_rec = mean_j P_rec F_c(y_j) (paper Eq. 5, 17).
+
+Auxiliary samples are drawn from a dedicated MS1M identity range disjoint from
+the proxy / defender training identities, mirroring the paper's use of a
+separate auxiliary set.
+
+Outputs checkpoints/bank.pt with aux_bank (M,C,56,56, fp16), mu_rec (C,56,56)
+and the fixed channel-shuffle permutation perm (C,).
+"""
+import os
+import glob
+import argparse
+
+import torch
+
+from models import load_ir50, ClientEncoder
+from decoyface import proj
+from data import load_rgb01
+from common import ADAFACE_CKPT, CKPT, CFG, DEVICE, fixed_perm
+
+AUX_ID_LO, AUX_ID_HI = 50000, 60000   # dedicated auxiliary identity range
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ms1m_root", required=True)
+    args = ap.parse_args()
+
+    dev = DEVICE
+    ir50 = load_ir50(ADAFACE_CKPT, dev)
+    Fc = ClientEncoder(ir50).to(dev).eval()
+    rssd = torch.load(CKPT / "rssd.pt", map_location=dev)
+    P_rec, P_N = rssd["P_rec"].to(dev), rssd["P_N"].to(dev)
+
+    ids = sorted(os.listdir(args.ms1m_root), key=int)[AUX_ID_LO:AUX_ID_HI]
+    paths = []
+    for ident in ids:
+        paths.extend(sorted(glob.glob(f"{args.ms1m_root}/{ident}/*.jpg")))
+    assert len(paths) >= CFG["M"], f"aux pool {len(paths)} < M {CFG['M']}"
+    g = torch.Generator().manual_seed(7)
+    pick = torch.randperm(len(paths), generator=g)[:CFG["M"]]
+    paths = [paths[i] for i in pick.tolist()]
+
+    bank = torch.empty(CFG["M"], CFG["C"], 56, 56, dtype=torch.float16)
+    mu_rec = torch.zeros(CFG["C"], 56, 56, device=dev)
+    bs = 128
+    with torch.no_grad():
+        for i in range(0, len(paths), bs):
+            imgs = torch.stack([load_rgb01(p) for p in paths[i:i + bs]]).to(dev)
+            z = Fc(imgs)
+            bank[i:i + imgs.shape[0]] = proj(P_N, z).half().cpu()
+            mu_rec += proj(P_rec, z).sum(0)
+    mu_rec /= len(paths)
+
+    torch.save({"aux_bank": bank, "mu_rec": mu_rec.cpu(),
+                "perm": fixed_perm(CFG["C"])}, CKPT / "bank.pt")
+    print(f"bank {tuple(bank.shape)} mu_rec {tuple(mu_rec.shape)} "
+          f"from {len(paths)} aux samples", flush=True)
+
+
+if __name__ == "__main__":
+    main()

build_rssd.py

@@ -0,0 +1,114 @@
+"""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()

checkpoints/adaface_ir50_ms1mv2.backbone.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:321a032dd99ce7ffeab078c92436e0371e8d752d66994c48825ff296cac58490
+size 174645926

checkpoints/bank.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba04ddb79804abb619daaae9eec4ca5c51155347d266422bd367505336851c6f
+size 2007845254

checkpoints/rssd.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d8629eaaa70ff55684262ce4e1f6c3b2be982d5c6becb36a3021ae260ecd8e2
+size 68024

common.py

@@ -0,0 +1,29 @@
+"""Shared config and paths for DecoyFace."""
+from pathlib import Path
+
+import torch
+
+CKPT = Path(__file__).resolve().parent / "checkpoints"
+ADAFACE_CKPT = CKPT / "adaface_ir50_ms1mv2.backbone.pth"
+
+# DecoyFace hyper-parameters (paper Sec. V-B "Implementation Details").
+CFG = {
+    "C": 64,
+    "r": 32,          # RSS rank
+    "G": 16,          # groups
+    "K": 3,           # preserved true-sample groups
+    "M": 5000,        # auxiliary complementary bank size
+    "sc_lo": 8, "sc_hi": 12,   # coherent subset size s_c ~ U[8,12]
+    "L_lo": 2, "L_hi": 4,      # shared source count L ~ U[2,4]
+    "t": 3,           # random scale exponent k ~ U[-3,3]
+    "q": 0.95,        # quantile
+    "tau_s": 0.8,     # server target magnitude
+    "eps": 1e-6,
+}
+
+DEVICE = "cuda:0"
+
+
+def fixed_perm(C=64, seed=1234):
+    g = torch.Generator().manual_seed(seed)
+    return torch.randperm(C, generator=g)

data.py

@@ -0,0 +1,35 @@
+"""MS1M dataset: identity folders of 112x112 ArcFace-aligned faces."""
+import os
+import glob
+
+import cv2
+import torch
+from torch.utils.data import Dataset
+
+
+def load_rgb01(path):
+    img = cv2.imread(path)              # BGR uint8, 112x112
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    t = torch.from_numpy(img).float().permute(2, 0, 1) / 255.0
+    return t
+
+
+class MS1MDataset(Dataset):
+    """Aligned MS1M faces. Uses `num_ids` identity folders (sorted
+    numerically) starting at `start` so different scripts see the same
+    deterministic split."""
+
+    def __init__(self, root, num_ids, start=0):
+        ids = sorted(os.listdir(root), key=int)[start:start + num_ids]
+        self.samples = []
+        for label, ident in enumerate(ids):
+            for f in sorted(glob.glob(f"{root}/{ident}/*.jpg")):
+                self.samples.append((f, label))
+        self.num_classes = len(ids)
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, i):
+        path, label = self.samples[i]
+        return load_rgb01(path), label

decoyface.py

@@ -0,0 +1,114 @@
+"""DecoyFace client-side protection (DGCM) and server-side recovery (ACM).
+
+All projections act on the channel dimension of a C x H x W feature at every
+spatial location. C = 64 (F_c output), r = 32, G = 16 groups of size 4.
+
+The RSSD statistics (P_rec, P_N, mu_rec), the fixed channel-shuffle permutation
+and the auxiliary complementary bank D_N are built once (see build_rssd.py /
+build_bank.py) and loaded into DecoyProtector.
+"""
+import torch
+import torch.nn.functional as F
+
+
+def proj(P, z):
+    """Apply channel-space projection P (C,C) to feature z (B,C,H,W)."""
+    b, c, h, w = z.shape
+    out = P @ z.reshape(b, c, h * w)
+    return out.reshape(b, c, h, w)
+
+
+def group_shuffle(z, perm, G):
+    """CSG: permute channels by `perm`, then split into G contiguous groups.
+    z (B,C,H,W) -> (B,G,c_g,H,W)."""
+    b, c, h, w = z.shape
+    z = z[:, perm]
+    return z.reshape(b, G, c // G, h, w)
+
+
+def group_unshuffle(zg, inv_perm):
+    """Inverse CSG: (B,G,c_g,H,W) -> (B,C,H,W) in the original channel order."""
+    b, G, cg, h, w = zg.shape
+    z = zg.reshape(b, G * cg, h, w)
+    return z[:, inv_perm]
+
+
+def group_instance_norm(zg, eps=1e-5):
+    """InstanceNorm applied to every group's channels over the spatial dims."""
+    b, G, cg, h, w = zg.shape
+    x = zg.reshape(b * G, cg, h, w)
+    x = F.instance_norm(x, eps=eps)
+    return x.reshape(b, G, cg, h, w)
+
+
+class DecoyProtector:
+    def __init__(self, P_rec, P_N, mu_rec, aux_bank, perm, cfg, device):
+        self.device = device
+        self.P_rec = P_rec.to(device)
+        self.P_N = P_N.to(device)
+        self.mu_rec = mu_rec.to(device)          # (C,H,W)
+        self.perm = perm.to(device)              # (C,)
+        self.inv_perm = torch.argsort(perm).to(device)
+        self.cfg = cfg
+        self.G = cfg["G"]
+        self.K = cfg["K"]
+        self.M = aux_bank.shape[0]
+        # Pre-group and pre-normalize the auxiliary complementary bank once.
+        # aux_bank holds n_j = P_N F_c(y_j); grouped+IN gives the decoy groups.
+        ag = group_shuffle(aux_bank.to(device), self.perm, self.G)
+        self.aux_grouped_norm = group_instance_norm(ag)   # (M,G,c_g,H,W)
+
+    def _sample_routing(self, B):
+        """Per-sample mask m (B,G) and decoy source indices src (B,G)."""
+        G, K, M, cfg = self.G, self.K, self.M, self.cfg
+        mask = torch.zeros(B, G, dtype=torch.bool)
+        src = torch.zeros(B, G, dtype=torch.long)
+        for b in range(B):
+            keep = torch.randperm(G)[:K]
+            mask[b, keep] = True
+            g0 = torch.nonzero(~mask[b], as_tuple=False).flatten()
+            sc = int(torch.randint(cfg["sc_lo"], cfg["sc_hi"] + 1, (1,)))
+            sc = min(sc, g0.numel())
+            L = int(torch.randint(cfg["L_lo"], cfg["L_hi"] + 1, (1,)))
+            L = min(L, sc) if sc > 0 else L
+            perm_g0 = g0[torch.randperm(g0.numel())]
+            coherent = perm_g0[:sc]
+            independent = perm_g0[sc:]
+            shared_pool = torch.randperm(M)[:L]
+            for g in coherent.tolist():
+                src[b, g] = shared_pool[torch.randint(L, (1,))]
+            for g in independent.tolist():
+                src[b, g] = torch.randint(M, (1,))
+        return mask.to(self.device), src.to(self.device)
+
+    def protect(self, z_A, z_B):
+        """Client-side DGCM. z_A, z_B: (B,C,H,W). Returns transmitted z_t."""
+        B = z_A.shape[0]
+        r_B = proj(self.P_rec, z_B)
+        n_A = proj(self.P_N, z_A)
+
+        uA = group_instance_norm(group_shuffle(n_A, self.perm, self.G))
+        mask, src = self._sample_routing(B)
+        ar = torch.arange(self.G, device=self.device)
+        decoy = self.aux_grouped_norm[src, ar]            # (B,G,c_g,H,W)
+        m = mask[:, :, None, None, None]
+        mixed = torch.where(m, uA, decoy)
+        n_tilde = group_unshuffle(mixed, self.inv_perm)
+
+        z_tilde = n_tilde + r_B
+        k = torch.randint(-self.cfg["t"], self.cfg["t"] + 1, (B, 1, 1, 1),
+                          device=self.device)
+        z_t = (2.0 ** k.float()) * z_tilde
+        return z_t
+
+    def server_recover(self, z_t):
+        """Server-side ACM. z_t: (B,C,H,W). Returns canonicalized z_can."""
+        z_s = proj(self.P_N, z_t)
+        zg = group_shuffle(z_s, self.perm, self.G)        # (B,G,c_g,H,W)
+        b, G, cg, h, w = zg.shape
+        flat = zg.reshape(b, G, cg * h * w).abs()
+        q = torch.quantile(flat, self.cfg["q"], dim=2)    # (B,G)
+        scale = self.cfg["tau_s"] / (q + self.cfg["eps"])
+        zg = zg * scale[:, :, None, None, None]
+        n_cal = group_unshuffle(zg, self.inv_perm)
+        return self.mu_rec[None] + n_cal

models.py

@@ -0,0 +1,182 @@
+"""Backbone (IR-50) split into F_c / F_s, and the U-Net feature-inversion network.
+
+The IR-50 definition follows face.evoLVe / AdaFace exactly so the pretrained
+AdaFace MS1Mv2 checkpoint loads without surgery. The client encoder F_c is the
+stem plus the first residual stage, whose output is 64x56x56 -- the intermediate
+feature used throughout DecoyFace (paper Sec. V-B). F_s is the remainder.
+"""
+from collections import namedtuple
+
+import torch
+import torch.nn as nn
+
+
+# ----------------------------- IR-50 backbone ------------------------------ #
+class Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+
+class BasicBlockIR(nn.Module):
+    def __init__(self, in_channel, depth, stride):
+        super().__init__()
+        if in_channel == depth:
+            self.shortcut_layer = nn.MaxPool2d(1, stride)
+        else:
+            self.shortcut_layer = nn.Sequential(
+                nn.Conv2d(in_channel, depth, (1, 1), stride, bias=False),
+                nn.BatchNorm2d(depth))
+        self.res_layer = nn.Sequential(
+            nn.BatchNorm2d(in_channel),
+            nn.Conv2d(in_channel, depth, (3, 3), (1, 1), 1, bias=False),
+            nn.BatchNorm2d(depth),
+            nn.PReLU(depth),
+            nn.Conv2d(depth, depth, (3, 3), stride, 1, bias=False),
+            nn.BatchNorm2d(depth))
+
+    def forward(self, x):
+        return self.res_layer(x) + self.shortcut_layer(x)
+
+
+Bottleneck = namedtuple('Block', ['in_channel', 'depth', 'stride'])
+
+
+def get_block(in_channel, depth, num_units, stride=2):
+    return [Bottleneck(in_channel, depth, stride)] + \
+           [Bottleneck(depth, depth, 1) for _ in range(num_units - 1)]
+
+
+def get_blocks_ir50():
+    return [
+        get_block(64, 64, 3),      # stage 1 -> 64 x 56 x 56
+        get_block(64, 128, 4),     # stage 2 -> 128 x 28 x 28
+        get_block(128, 256, 14),   # stage 3 -> 256 x 14 x 14
+        get_block(256, 512, 3),    # stage 4 -> 512 x 7 x 7
+    ]
+
+
+class IR50(nn.Module):
+    """Full IR-50 recognizer (input 112x112 -> 512-d embedding)."""
+
+    def __init__(self):
+        super().__init__()
+        self.input_layer = nn.Sequential(
+            nn.Conv2d(3, 64, (3, 3), 1, 1, bias=False),
+            nn.BatchNorm2d(64), nn.PReLU(64))
+        modules = []
+        for block in get_blocks_ir50():
+            for b in block:
+                modules.append(BasicBlockIR(b.in_channel, b.depth, b.stride))
+        self.body = nn.Sequential(*modules)
+        self.output_layer = nn.Sequential(
+            nn.BatchNorm2d(512), nn.Dropout(0.4), Flatten(),
+            nn.Linear(512 * 7 * 7, 512), nn.BatchNorm1d(512, affine=False))
+
+    def forward(self, x):
+        x = self.input_layer(x)
+        x = self.body(x)
+        return self.output_layer(x)
+
+
+# The first residual stage (stage 1) spans body[0:3]; splitting there yields the
+# 64x56x56 intermediate feature the paper transmits.
+SPLIT_INDEX = 3
+
+
+def adaface_preprocess(x_rgb01):
+    """RGB image in [0,1] -> AdaFace input (BGR, normalized to [-1,1])."""
+    x_bgr = x_rgb01.flip(1)  # RGB -> BGR on the channel dim
+    return (x_bgr - 0.5) / 0.5
+
+
+class ClientEncoder(nn.Module):
+    """F_c: AdaFace stem + stage 1. Frozen. Consumes RGB [0,1], emits 64x56x56."""
+
+    def __init__(self, ir50: IR50):
+        super().__init__()
+        self.input_layer = ir50.input_layer
+        self.stage1 = ir50.body[:SPLIT_INDEX]
+
+    def forward(self, x_rgb01):
+        x = adaface_preprocess(x_rgb01)
+        x = self.input_layer(x)
+        return self.stage1(x)
+
+
+class ServerBackbone(nn.Module):
+    """F_s: AdaFace stage 2-4 + output layer. 64x56x56 -> 512-d embedding."""
+
+    def __init__(self, ir50: IR50):
+        super().__init__()
+        self.rest = ir50.body[SPLIT_INDEX:]
+        self.output_layer = ir50.output_layer
+
+    def forward(self, feat):
+        return self.output_layer(self.rest(feat))
+
+
+def load_ir50(ckpt_path, device):
+    net = IR50()
+    sd = torch.load(ckpt_path, map_location="cpu")
+    if isinstance(sd, dict) and "state_dict" in sd:
+        sd = sd["state_dict"]
+    missing, unexpected = net.load_state_dict(sd, strict=True)
+    net.eval().to(device)
+    return net
+
+
+# ------------------------------- U-Net ------------------------------------- #
+class ConvBlock(nn.Module):
+    def __init__(self, cin, cout, norm="batch"):
+        super().__init__()
+        if norm == "batch":
+            norm1 = nn.BatchNorm2d(cout)
+            norm2 = nn.BatchNorm2d(cout)
+        elif norm == "instance":
+            norm1 = nn.InstanceNorm2d(cout)
+            norm2 = nn.InstanceNorm2d(cout)
+        else:
+            raise ValueError(norm)
+        self.net = nn.Sequential(
+            nn.Conv2d(cin, cout, 3, 1, 1, bias=False),
+            norm1, nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv2d(cout, cout, 3, 1, 1, bias=False),
+            norm2, nn.LeakyReLU(0.2, inplace=True))
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class UNetInverter(nn.Module):
+    """Feature-inversion U-Net: 64x56x56 feature -> 3x112x112 RGB image in [0,1].
+
+    Encoder-decoder with skip connections over the feature map, plus a final
+    upsampling stage to reach the 112x112 image resolution.
+    """
+
+    def __init__(self, in_ch=64, norm="batch"):
+        super().__init__()
+        self.enc0 = ConvBlock(in_ch, 64, norm)    # 56
+        self.down1 = nn.MaxPool2d(2)              # 28
+        self.enc1 = ConvBlock(64, 128, norm)
+        self.down2 = nn.MaxPool2d(2)              # 14
+        self.enc2 = ConvBlock(128, 256, norm)
+        self.bott = ConvBlock(256, 256, norm)
+
+        self.up2 = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)
+        self.dec2 = ConvBlock(256 + 128, 128, norm)  # 28
+        self.up1 = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)
+        self.dec1 = ConvBlock(128 + 64, 64, norm)  # 56
+        self.up0 = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False)
+        self.dec0 = ConvBlock(64, 32, norm)       # 112
+        self.head = nn.Conv2d(32, 3, 3, 1, 1)
+
+    def forward(self, feat):
+        e0 = self.enc0(feat)
+        e1 = self.enc1(self.down1(e0))
+        e2 = self.enc2(self.down2(e1))
+        b = self.bott(e2)
+        d2 = self.dec2(torch.cat([self.up2(b), e1], 1))
+        d1 = self.dec1(torch.cat([self.up1(d2), e0], 1))
+        d0 = self.dec0(self.up0(d1))
+        return torch.sigmoid(self.head(d0))

train_defender.py

@@ -0,0 +1,111 @@
+"""Defender training (paper Sec. IV-E): F_c frozen; train F_s + ArcFace head on
+the server-side canonicalized feature z_can with V=2 protected views and the
+multi-view consistency loss.  L = L_id + lambda_cons * L_cons.
+
+Outputs checkpoints/defender.pt with the trained F_s and ArcFace head.
+"""
+import copy
+import argparse
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+import wandb
+
+from models import load_ir50, ClientEncoder, ServerBackbone
+from data import MS1MDataset
+from decoyface import DecoyProtector
+from common import ADAFACE_CKPT, CKPT, CFG, DEVICE
+
+V = 2
+LAMBDA_CONS = 0.1
+BATCH = 256
+EPOCHS = 15
+LR = 0.1
+DECAY_EPOCHS = [4, 8, 12]
+
+
+class ArcFace(nn.Module):
+    def __init__(self, in_dim, n_cls, s=64.0, m=0.5):
+        super().__init__()
+        self.W = nn.Parameter(torch.empty(n_cls, in_dim))
+        nn.init.xavier_normal_(self.W)
+        self.s, self.m = s, m
+
+    def forward(self, emb, y):
+        emb = F.normalize(emb)
+        w = F.normalize(self.W)
+        cos = emb @ w.T
+        theta = torch.acos(cos.clamp(-1 + 1e-6, 1 - 1e-6))
+        target = torch.cos(theta + self.m)
+        onehot = F.one_hot(y, cos.shape[1]).float()
+        logits = self.s * (onehot * target + (1 - onehot) * cos)
+        return F.cross_entropy(logits, y)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ms1m_root", required=True)
+    ap.add_argument("--num_ids", type=int, default=30000)
+    ap.add_argument("--epochs", type=int, default=EPOCHS)
+    args = ap.parse_args()
+    dev = DEVICE
+
+    ir50 = load_ir50(ADAFACE_CKPT, dev)
+    Fc = ClientEncoder(ir50).to(dev).eval()
+    for p in Fc.parameters():
+        p.requires_grad_(False)
+    Fs = ServerBackbone(copy.deepcopy(ir50)).to(dev).train()
+    rssd = torch.load(CKPT / "rssd.pt", map_location=dev)
+    bank = torch.load(CKPT / "bank.pt", map_location=dev)
+    protector = DecoyProtector(rssd["P_rec"], rssd["P_N"], bank["mu_rec"],
+                               bank["aux_bank"].float(), bank["perm"], CFG, dev)
+
+    ds = MS1MDataset(args.ms1m_root, args.num_ids)
+    head = ArcFace(512, ds.num_classes).to(dev)
+    loader = DataLoader(ds, batch_size=BATCH, shuffle=True, num_workers=12,
+                        drop_last=True, pin_memory=True, persistent_workers=True)
+
+    params = list(Fs.parameters()) + list(head.parameters())
+    opt = torch.optim.SGD(params, lr=LR, momentum=0.9, weight_decay=5e-4)
+    wandb.init(project="decoyface-repro", name="defender", mode="online",
+               config={**CFG, "num_ids": args.num_ids, "epochs": args.epochs})
+
+    step = 0
+    for epoch in range(args.epochs):
+        lr = LR * (0.1 ** sum(epoch >= d for d in DECAY_EPOCHS))
+        for g in opt.param_groups:
+            g["lr"] = lr
+        for x, y in loader:
+            x, y = x.to(dev, non_blocking=True), y.to(dev, non_blocking=True)
+            with torch.no_grad():
+                z = Fc(x)
+            embs = []
+            id_loss = 0.0
+            for _ in range(V):
+                z_B = z.roll(torch.randint(1, x.shape[0], (1,)).item(), 0)
+                z_can = protector.server_recover(protector.protect(z, z_B))
+                e = Fs(z_can)
+                id_loss = id_loss + head(e, y)
+                embs.append(F.normalize(e))
+            id_loss = id_loss / V
+            cons = ((embs[0] - embs[1]) ** 2).sum(1).mean()
+            loss = id_loss + LAMBDA_CONS * cons
+            opt.zero_grad(set_to_none=True)
+            loss.backward()
+            opt.step()
+            if step % 100 == 0:
+                wandb.log({"def/loss": loss.item(), "def/id": id_loss.item(),
+                           "def/cons": cons.item(), "def/lr": lr}, step=step)
+            if step % 1000 == 0:
+                print(f"ep{epoch} step{step} loss{loss.item():.3f} "
+                      f"id{id_loss.item():.3f} cons{cons.item():.4f}", flush=True)
+            step += 1
+        torch.save({"Fs": Fs.state_dict(), "head": head.state_dict()},
+                   CKPT / "defender.pt")
+    wandb.finish()
+
+
+if __name__ == "__main__":
+    main()