| """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() |
|
|
