File size: 4,611 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 | """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
|