File size: 4,105 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 | """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()
|