emilyxuan/accesscontrol

1import random2import numpy as np3import torch4import matplotlib.pyplot as plt5from sklearn.decomposition import PCA6 7from trigger_encoder import TriggerEncoder8 9 10def build_unseen_trigger(img_size, num_points=2, grid_size=8, seed=999, square_size=3):11    assert square_size >= 1 and square_size % 2 == 112 13    margin = max(2, img_size // 8)14    xs = torch.linspace(margin, img_size - margin - 1, steps=grid_size).long()15    ys = torch.linspace(margin, img_size - margin - 1, steps=grid_size).long()16    candidate_positions = [(int(y), int(x)) for y in ys for x in xs]17 18    trig = torch.full((3, img_size, img_size), -1.0, dtype=torch.float32)19    rng = random.Random(int(seed))20    selected = rng.sample(candidate_positions, num_points)21 22    half = square_size // 223    for (y, x) in selected:24        y1 = max(0, y - half)25        y2 = min(img_size, y + half + 1)26        x1 = max(0, x - half)27        x2 = min(img_size, x + half + 1)28        trig[:, y1:y2, x1:x2] = 1.029 30    return trig31 32 33def main():34    ckpt_path = "./logs/cifar10_cond_acsctrl/cond128_neg0.25_0.25_seed42/cond128_neg0.5_seed42_20260426_003254/checkpoints/global_ckpt_round240.pt"35    device = "cuda" if torch.cuda.is_available() else "cpu"36 37    ckpt = torch.load(ckpt_path, map_location="cpu")38    cfg = ckpt["config"]39 40    img_size = cfg["img_size"]41    cond_dim = cfg["cond_dim"]42    trigger_base_ch = cfg["trigger_base_ch"]43 44    encoder = TriggerEncoder(45        in_ch=3,46        base_ch=trigger_base_ch,47        cond_dim=cond_dim,48        img_size=img_size,49        grid_size=8,50        hidden_dim=256,51        use_coord_feat=True,52    ).to(device)53    encoder.load_state_dict(ckpt["trigger_encoder"], strict=True)54    encoder.eval()55 56    client_triggers = [t.clone() for t in ckpt["client_triggers"]]57 58    xs = []59    labels = []60 61    with torch.no_grad():62        # trained triggers63        for i, trig in enumerate(client_triggers):64            emb = encoder(trig.unsqueeze(0).to(device)).squeeze(0).cpu().numpy()65            xs.append(emb)66            labels.append(f"trained_{i}")67 68        # fixed unseen triggers69        for s in [999, 1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010, 1011, 1012, 1013, 1014, 1015, 1016, 1017, 1018, 1019]:70            trig = build_unseen_trigger(71                img_size=img_size,72                num_points=cfg.get("unseen_num_points", 2),73                grid_size=cfg.get("unseen_grid_size", 8),74                seed=s,75                square_size=cfg.get("unseen_square_size", 3),76            )77            emb = encoder(trig.unsqueeze(0).to(device)).squeeze(0).cpu().numpy()78            xs.append(emb)79            labels.append(f"unseen_{s}")80 81        # none82        none_emb = np.zeros(cond_dim, dtype=np.float32)83        xs.append(none_emb)84        labels.append("none")85 86    X = np.stack(xs, axis=0)87    pca = PCA(n_components=2)88    X2 = pca.fit_transform(X)89 90    plt.figure(figsize=(8, 6))91    for i, name in enumerate(labels):92        if name.startswith("trained_"):93            marker = "o"94        elif name.startswith("unseen_"):95            marker = "^"96        else:97            marker = "x"98 99        plt.scatter(X2[i, 0], X2[i, 1], marker=marker)100        plt.text(X2[i, 0] + 0.02, X2[i, 1] + 0.02, name, fontsize=8)101 102    plt.xlabel(f"PC1 ({pca.explained_variance_ratio_[0] * 100:.1f}%)")103    plt.ylabel(f"PC2 ({pca.explained_variance_ratio_[1] * 100:.1f}%)")104    plt.title("Trigger Embedding PCA")105    plt.tight_layout()106    plt.savefig("trigger_pca240.png", dpi=200)107    print("saved: trigger_pca220.png")108 109 110if __name__ == "__main__":111    main()