QuadAttack release

.gitignore

@@ -0,0 +1,143 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+
+*.tar.gz
+cifar-10-batches-py
+logs/**/**
+logs_old/**/**
+config-lock.yaml
+
+*.png
+*.p
+datasets/*
+data/*
+*.save
+*.txt

README.md

@@ -0,0 +1,11 @@
+## Getting Started
+
+Install the dependencies with `pip`:
+pip install -r requirements.txt 
+
+python setup.py develop
+
+Adjust configuration in base_config.yaml
+
+Run current configuration by
+python modelguidedattacks/run.py base_config.yaml

base_config.yaml

@@ -0,0 +1,36 @@
+seed: 10
+nproc_per_node: 1
+k: 5
+data_path: ./
+train_batch_size: 64
+eval_batch_size: 64
+num_workers: 4
+max_epochs: 1000
+train_epoch_length: null
+eval_epoch_length: null
+lr: 0.001
+unguided_lr: 0.0022
+use_amp: false
+debug: false
+model: resnet50
+dataset: imagenet
+output_dir: ./logs
+log_every_iters: 1
+unguided_iterations: 30
+overfit: false
+guide_model: "unguided"
+loss: "cvxproj"
+out_dir: ""
+attack_sampling: "random"
+
+cvx_proj_margin: 0.2
+topk_loss_coef_upper: 20
+opt_warmup_its: 5
+binary_search_steps: 1
+
+dump_plots: false
+plot_idx: find # or specific batch idx
+plot_out: "myplots/mymethod"
+
+# List of models used to find correct subset
+compare_models: ["resnet50", "deit_small", "vit_base", "densenet121"]

debug_test.py

@@ -0,0 +1,185 @@
+import cvxpy as cp
+from cvxpylayers.torch import CvxpyLayer
+
+from torch.nn import functional as F
+import torch
+from modelguidedattacks import cls_models
+import time
+
+torch.manual_seed(0)
+device = "cuda"
+# model = cls_models.get_model("imagenet", "resnet18", device)
+
+rand_feats = torch.randn(1, 512, device=device)
+attack_targets = [4, 7, 5, 9, 2]
+
+# # pred_logits = model.head(rand_feats)
+
+# # head_W, head_bias = model.head_matrices()
+
+(head_W, head_bias, pred_logits) = torch.load("debugsaveimagenet.save")
+
+rand_feats, rand_logits, attack_targets = torch.load("attack_case.p", map_location=device)
+reconstructed_logits = rand_feats@head_W.T + head_bias
+
+num_feats = head_W.shape[1]
+num_classes = head_W.shape[0]
+x = cp.Variable(num_feats)
+
+anchor_feats = cp.Parameter(x.shape)
+A = cp.Parameter(head_W.shape)
+b = cp.Parameter(head_bias.shape)
+
+logits = A@x + b
+
+MARGIN = 0.1
+
+# constraints  = []
+# for i in range(len(attack_targets) - 1):
+#     constraints.append( logits[attack_targets[i]] - logits[attack_targets[i+1]] >= MARGIN)
+
+# for i in range(num_classes):
+#     if i in attack_targets:
+#         continue
+
+#     constraints.append(logits[attack_targets[-1]] - logits[i] >= MARGIN )
+
+# objective = cp.Minimize(0.5 * cp.pnorm(x - anchor_feats, p=2))
+# problem = cp.Problem(objective, constraints)
+
+# anchor_feats.value = rand_feats[0].cpu().numpy()
+# A.value = head_W.detach().cpu().numpy()
+# b.value = head_bias.detach().cpu().numpy()
+
+# start_time = time.time()
+# problem.solve()
+# print ("Non vectorized sol", time.time() - start_time)
+
+# logits_sol_torch = torch.from_numpy(logits.value)
+# logits_check = logits_sol_torch.argsort(descending=True)
+
+# feats_sol = torch.from_numpy(x.value[:, None]).float().to(rand_feats)
+# sol_feat_norm = (feats_sol[:, 0].cpu() - rand_feats[0].cpu()).norm(dim=-1)
+# sol_logits = head_W@feats_sol + head_bias[:, None]
+# sol_sort = sol_logits.argsort(dim=0, descending=True)
+
+
+# Constraint matrix
+num_constraints = num_classes - 1
+D = torch.zeros((num_classes), num_constraints)
+
+non_attack_targets = list(set(range(num_classes)) - set(attack_targets))
+
+for constraint_cursor in range(num_constraints):
+    if constraint_cursor < len(attack_targets) - 1:
+        D[attack_targets[constraint_cursor], constraint_cursor] = 1
+        D[attack_targets[constraint_cursor + 1], constraint_cursor] = -1
+    else:
+        non_attack_i = constraint_cursor - len(attack_targets) + 1
+        D[attack_targets[-1], constraint_cursor] = 1
+        D[non_attack_targets[non_attack_i], constraint_cursor] = -1
+
+D = D.T
+# vectorized_differences = D @ logits
+# vectorized_constraint = vectorized_differences >= torch.full(vectorized_differences.shape, fill_value=MARGIN).numpy()
+
+# Q = 2*torch.eye(x.shape[0]).numpy()
+# P = -2*anchor_feats
+
+# G = D@A
+# H = MARGIN - D @ b
+
+# G = -G
+# H = -H
+
+# vectorized_constraint = G@x <= H
+
+# objective = cp.Minimize((1/2)*cp.quad_form(x, Q) + P.T@x)
+# problem = cp.Problem(objective, [vectorized_constraint])
+
+# anchor_feats.value = rand_feats[0].cpu().numpy()
+# A.value = head_W.detach().cpu().numpy()
+# b.value = head_bias.detach().cpu().numpy()
+
+# start_time = time.time()
+# problem.solve()
+# print ("vectorized sol", time.time() - start_time)
+
+# logits_sol_torch = torch.from_numpy(logits.value)
+# logits_check = logits_sol_torch.argsort(descending=True)
+# feats_sol = torch.from_numpy(x.value[:, None]).float().to(rand_feats)
+# sol_feat_norm = (feats_sol[:, 0].cpu() - rand_feats[0].cpu()).norm(dim=-1)
+# sol_logits = head_W@feats_sol + head_bias[:, None]
+# sol_sort = sol_logits.argsort(dim=0, descending=True)
+
+import qpth
+
+
+B = 2
+nz = num_feats
+nineq = num_constraints
+device = "cuda"
+
+attack_targets = attack_targets.expand(B, -1)
+K = attack_targets.shape[-1]
+
+# Start with all classes should be less than smallest attack target
+D = -torch.eye(num_classes, device=device)[None].repeat(B, 1, 1)
+attack_targets_write = attack_targets[:, -1][:, None, None].expand(-1, D.shape[1], -1)
+D.scatter_(dim=2, index=attack_targets_write, src=torch.ones(attack_targets_write.shape, device=device))
+
+# Clear out the constraint row for each item in the attack targets
+attack_targets_clear = attack_targets[:, :, None].expand(-1, -1, D.shape[-1])
+D.scatter_(dim=1, index=attack_targets_clear, src=torch.zeros(attack_targets_clear.shape, device=device))
+
+batch_inds = torch.arange(B, device=device)[:, None].expand(-1, K - 1)
+attack_targets_pos = attack_targets[:, :-1] # [B, K-1]
+attack_targets_neg = attack_targets[:, 1:] # [B, K-1]
+
+attack_targets_neg_inds = torch.stack((
+    batch_inds,
+    attack_targets_neg,
+    attack_targets_neg
+), dim=0) # [3, B, K - 1]
+attack_targets_neg_inds = attack_targets_neg_inds.view(3, -1)
+
+D[attack_targets_neg_inds[0], attack_targets_neg_inds[1], attack_targets_neg_inds[2]] = -1
+
+attack_targets_pos_inds = torch.stack((
+    batch_inds,
+    attack_targets_neg,
+    attack_targets_pos
+), dim=0) # [3, B, K - 1]
+
+D[attack_targets_pos_inds[0], attack_targets_pos_inds[1], attack_targets_pos_inds[2]] = 1
+
+A = head_W.detach().to(device)
+b = head_bias.detach().to(device)
+D = D.to(device)
+
+#rand_feats: [B, num_features]
+Q = 2*torch.eye(nz, device=device)[None].expand(B, -1, -1)
+P = -2*rand_feats.to(device).expand(B, -1)
+
+# G = torch.randn(B, nineq, nz, device=device)
+G = -D@A
+
+# h = torch.randn(B, nineq)
+H = -(MARGIN - D @ b)
+
+# Constraints are indexed by smaller logit
+# First attack target isn't smaller than any logit, so its 
+# constraint index is redundant, but we keep it for easier parallelization
+# Make this constraint all 0s
+zero_inds = attack_targets[:, 0:1] # [B, 1]
+H.scatter_(dim=1, index=zero_inds, src=torch.zeros(zero_inds.shape, device=device))
+
+e = torch.empty(0, device=device)
+
+Q_t, P_t, G_t, H_t = torch.load("qpinputs.p", map_location=device)
+
+z_sol = qpth.qp.QPFunction(verbose=True, check_Q_spd=False)(Q, P, G, H, e, e).T
+
+logits = A@z_sol + b[:, None]
+
+x = 5

modelguidedattacks/cls_models/__init__.py

@@ -0,0 +1,4 @@
+from .registry import register_default_models, get_model
+from .accuracy import get_correct_subset, get_correct_subset_for_models
+
+register_default_models()

modelguidedattacks/cls_models/accuracy.py

@@ -0,0 +1,115 @@
+import logging
+import os
+
+import torch
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from modelguidedattacks.data import get_dataset
+from . import get_model
+
+from .registry import ClsModel
+from typing import Optional, List
+
+DATASET_METADATA_DIR = "./dataset_metadata"
+
+def correct_subset_cache_path(dataset_name: str, model_name: str, train: bool):
+    filename_train_val = "train" if train else "val"
+    subset_cache_filename = f"{dataset_name}_{model_name}_{filename_train_val}.p"
+    subset_cache_path = os.path.join(DATASET_METADATA_DIR, subset_cache_filename)
+
+    return subset_cache_path
+
+@torch.no_grad()
+def get_correct_subset(model: Optional[ClsModel]=None, dataset_name: Optional[str]=None, 
+                       model_name: Optional[str]=None, train=True, batch_size=256, 
+                       force_cache=False, device="cuda"):
+    """
+    model: Model to evaluate
+    dataset_name: Name of dataset (not needed if model is provided)
+    model_name: Name of model (not needed if model is provided)
+    train: Use training dataset
+    batch_size: Batch size to use while evaluating
+    force_cache: Only read from cache and fail if not available
+
+    Returns indices in dataset of correctly classified items
+    """
+
+    if model is not None:
+        assert dataset_name is None
+        assert model_name is None
+
+    if dataset_name is not None or model_name is not None:
+        assert dataset_name is not None
+        assert model_name is not None
+        assert model is None
+
+    if dataset_name is None:
+        dataset_name = model.dataset_name
+
+    if model_name is None:
+        model_name = model.model_name
+
+    filename_train_val = "train" if train else "val"
+    subset_cache_filename = f"{dataset_name}_{model_name}_{filename_train_val}.p"
+    subset_cache_path = os.path.join(DATASET_METADATA_DIR, subset_cache_filename)
+
+    os.makedirs(DATASET_METADATA_DIR, exist_ok=True)
+
+    if os.path.exists(subset_cache_path):
+        correct_subset = torch.load(subset_cache_path)
+        return correct_subset
+
+    if force_cache:
+        raise Exception("Cache not found and requested for cached correct subset.")
+
+    logging.info(f"No cache found. Computing correct subset for {dataset_name}-{model_name} Train: {train}")
+
+    device = device if model is None else model.device    
+
+    if model is None:
+        model = get_model(dataset_name, model_name, device)
+    
+    model.eval()
+
+    train_dataset, val_dataset = get_dataset(dataset_name)
+
+    dataset = train_dataset
+    
+    if not train:
+        dataset = val_dataset
+
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+
+    correct_indices = []
+
+    for batch_i, (batch_imgs, batch_gt_class) in tqdm(enumerate(dataloader), total=len(dataloader)):
+        if torch.device(model.device).type.startswith("cuda"):
+            torch.cuda.synchronize(model.device)
+
+        data_start_index = batch_i * batch_size
+        predictions = model(batch_imgs.to(model.device)) # [B, C]
+        prediction_class_idx = predictions.argmax(dim=-1) # [B] (long)
+        prediction_correct = prediction_class_idx == batch_gt_class.to(model.device)
+        batch_correct_idxs = data_start_index + prediction_correct.nonzero()[:, 0]
+        batch_correct_idxs = batch_correct_idxs.tolist()
+
+        correct_indices.extend(batch_correct_idxs)
+
+    correct_subset = set(correct_indices)
+    torch.save(correct_subset, subset_cache_path)
+
+    return set(correct_indices)
+
+def get_correct_subset_for_models(model_names: List[str], dataset_name, device, train):
+    correct_intersection = None
+    for model_name in model_names:
+        model_correct_subset = get_correct_subset(model_name=model_name, dataset_name=dataset_name,
+                                                   device=device, train=train)
+        
+        if correct_intersection is None:
+            correct_intersection = model_correct_subset
+        else:
+            correct_intersection = model_correct_subset.intersection(correct_intersection)
+
+    return list(correct_intersection)

modelguidedattacks/cls_models/registry.py

@@ -0,0 +1,163 @@
+import mmpretrain
+import torch
+from torch import nn
+from collections.abc import Iterable
+from mmpretrain.models.utils.attention import MultiheadAttention
+# This holds model instantiation functions by (dataset_name, model_name) tuple keys
+MODEL_REGISTRY = {}
+
+class ClsModel(nn.Module):
+    dataset_name: str
+    model_name: str
+
+    def __init__(self, dataset_name: str, model_name: str, device: str) -> None:
+        super().__init__()
+        self.dataset_name = dataset_name
+        self.model_name = model_name
+        self.device = device
+
+    def head_features(self):
+        pass
+
+    def num_classes(self):
+        pass
+
+    def forward(self, x):
+        """
+        x: [B, 3 (RGB), H, W] image (float) [0,1]
+
+        returns: [B, C] class logits
+        """
+
+        raise NotImplementedError("Forward not implemented for base class")
+
+class MMPretrainModelWrapper(ClsModel):
+    """
+    Calls data preprocessing for model before entering forward
+    """
+    def __init__(self, model: nn.Module, dataset_name: str, model_name: str, device: str) -> None:
+        super().__init__(dataset_name, model_name, device)
+        self.model = model
+
+    @property
+    def final_linear_layer(self):
+        return self.model.head.fc
+
+    def head_features(self):
+        return self.final_linear_layer.in_features
+    
+    def num_classes(self):
+        return self.final_linear_layer.out_features
+    
+    def head(self, feats):
+        return self.model.head((feats,))
+    
+    def head_matrices(self):
+        return self.final_linear_layer.weight, self.final_linear_layer.bias
+
+    def forward(self, x, return_features=False):
+        # Data preprocessor expects 0-255 range, but we don't want to cast to proper
+        # uint8 because we want to maintain differentiability
+        x = x * 255.
+        x = self.model.data_preprocessor({"inputs": x})["inputs"]
+
+        if return_features:
+            feats = self.model.extract_feat(x)
+            
+            preds = self.model.head(feats)
+            if isinstance(feats, Iterable):
+                feats = feats[-1]
+                
+            return preds, feats
+        else:
+            return self.model(x)
+        
+class MMPretrainVisualTransformerWrapper(MMPretrainModelWrapper):
+    def __init__(self, model, dataset_name: str, model_name: str, device: str) -> None:
+        super().__init__(model, dataset_name, model_name, device)
+
+        attn_layers = []
+
+        def find_mha(m: nn.Module):
+            if isinstance(m, MultiheadAttention):
+                attn_layers.append(m)
+
+        model.apply(find_mha)
+
+        self.attn_layers = attn_layers
+
+    @property
+    def final_linear_layer(self):
+        return self.model.head.layers.head
+    
+    def get_attention_maps(self, x):
+        clean_forwards = []
+
+        attention_maps = []
+
+        for attn_layer in self.attn_layers:
+            clean_forward = attn_layer.forward
+            clean_forwards.append(clean_forward)
+
+            def scaled_dot_prod_attn(query,
+                                        key,
+                                        value,
+                                        attn_mask=None,
+                                        dropout_p=0.,
+                                        scale=None,
+                                        is_causal=False):
+                scale = scale or query.size(-1)**0.5
+                if is_causal and attn_mask is not None:
+                    attn_mask = torch.ones(
+                        query.size(-2), key.size(-2), dtype=torch.bool).tril(diagonal=0)
+                if attn_mask is not None and attn_mask.dtype == torch.bool:
+                    attn_mask = attn_mask.masked_fill(not attn_mask, -float('inf'))
+
+                attn_weight = query @ key.transpose(-2, -1) / scale
+                if attn_mask is not None:
+                    attn_weight += attn_mask
+                attn_weight = torch.softmax(attn_weight, dim=-1)
+
+                attention_maps.append(attn_weight)
+
+                attn_weight = torch.dropout(attn_weight, dropout_p, True)
+                return attn_weight @ value
+
+            attn_layer.scaled_dot_product_attention = scaled_dot_prod_attn
+
+        ret_val = super().forward(x, False)
+
+        for attn_layer, clean_forward in zip(self.attn_layers, clean_forwards):
+            attn_layer.forward = clean_forward
+
+        return attention_maps
+    
+def register_mmcls_model(config_name, dataset_name, model_name, 
+                         wrapper_class=MMPretrainModelWrapper):
+    def instantiate_model(device):
+        model = mmpretrain.get_model(config_name, pretrained=True, device=device)
+        wrapper = wrapper_class(model, dataset_name, model_name, device)
+        return wrapper
+    
+    MODEL_REGISTRY[(dataset_name, model_name)] = instantiate_model
+
+def register_default_models():
+    register_mmcls_model("resnet18_8xb16_cifar10", "cifar10", "resnet18")
+    register_mmcls_model("resnet34_8xb16_cifar10", "cifar10", "resnet34")
+    register_mmcls_model("resnet18_8xb32_in1k", "imagenet", "resnet18")
+    register_mmcls_model("resnet50_8xb16_cifar100", "cifar100", "resnet50")
+    register_mmcls_model("resnet50_8xb32_in1k", "imagenet", "resnet50")
+    register_mmcls_model("densenet121_3rdparty_in1k", "imagenet", "densenet121")
+
+    register_mmcls_model("deit-small_4xb256_in1k", "imagenet", "deit_small",
+                          wrapper_class=MMPretrainVisualTransformerWrapper)
+    
+    register_mmcls_model("vit-base-p16_32xb128-mae_in1k", "imagenet", "vit_base",
+                          wrapper_class=MMPretrainVisualTransformerWrapper)
+
+def get_model(dataset_name, model_name, device):
+    """
+    Returns instance of model pretrained with specified dataset
+    """
+
+    return MODEL_REGISTRY[(dataset_name, model_name)](device).eval()

modelguidedattacks/data/__init__.py

@@ -0,0 +1,4 @@
+from .registry import register_default_datasets
+from .registry import get_dataset
+
+register_default_datasets()

modelguidedattacks/data/classification_wrapper.py

@@ -0,0 +1,22 @@
+import torch
+import torch.utils.data as data
+
+class TopKClassificationWrapper(data.Dataset):
+    def __init__(self, dataset: data.Dataset, attack_labels, seed=0, k=1) -> None:
+        super().__init__()
+        self.generator = torch.Generator("cpu")
+        self.generator.manual_seed(seed)
+
+        # Pregenerate attack labels
+        num_classes = len(dataset.classes)
+
+        self.src_dataset = dataset
+        self.attack_labels = attack_labels
+
+    def __getitem__(self, index):
+        image, label = self.src_dataset[index]
+
+        return image, label, self.attack_labels[index], index
+    
+    def __len__(self):
+        return len(self.src_dataset)

modelguidedattacks/data/imagenet_metadata.py

@@ -0,0 +1,1000 @@
+imgnet_idx_to_name = {0: 'tench, Tinca tinca',
+ 1: 'goldfish, Carassius auratus',
+ 2: 'great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias',
+ 3: 'tiger shark, Galeocerdo cuvieri',
+ 4: 'hammerhead, hammerhead shark',
+ 5: 'electric ray, crampfish, numbfish, torpedo',
+ 6: 'stingray',
+ 7: 'cock',
+ 8: 'hen',
+ 9: 'ostrich, Struthio camelus',
+ 10: 'brambling, Fringilla montifringilla',
+ 11: 'goldfinch, Carduelis carduelis',
+ 12: 'house finch, linnet, Carpodacus mexicanus',
+ 13: 'junco, snowbird',
+ 14: 'indigo bunting, indigo finch, indigo bird, Passerina cyanea',
+ 15: 'robin, American robin, Turdus migratorius',
+ 16: 'bulbul',
+ 17: 'jay',
+ 18: 'magpie',
+ 19: 'chickadee',
+ 20: 'water ouzel, dipper',
+ 21: 'kite',
+ 22: 'bald eagle, American eagle, Haliaeetus leucocephalus',
+ 23: 'vulture',
+ 24: 'great grey owl, great gray owl, Strix nebulosa',
+ 25: 'European fire salamander, Salamandra salamandra',
+ 26: 'common newt, Triturus vulgaris',
+ 27: 'eft',
+ 28: 'spotted salamander, Ambystoma maculatum',
+ 29: 'axolotl, mud puppy, Ambystoma mexicanum',
+ 30: 'bullfrog, Rana catesbeiana',
+ 31: 'tree frog, tree-frog',
+ 32: 'tailed frog, bell toad, ribbed toad, tailed toad, Ascaphus trui',
+ 33: 'loggerhead, loggerhead turtle, Caretta caretta',
+ 34: 'leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea',
+ 35: 'mud turtle',
+ 36: 'terrapin',
+ 37: 'box turtle, box tortoise',
+ 38: 'banded gecko',
+ 39: 'common iguana, iguana, Iguana iguana',
+ 40: 'American chameleon, anole, Anolis carolinensis',
+ 41: 'whiptail, whiptail lizard',
+ 42: 'agama',
+ 43: 'frilled lizard, Chlamydosaurus kingi',
+ 44: 'alligator lizard',
+ 45: 'Gila monster, Heloderma suspectum',
+ 46: 'green lizard, Lacerta viridis',
+ 47: 'African chameleon, Chamaeleo chamaeleon',
+ 48: 'Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis',
+ 49: 'African crocodile, Nile crocodile, Crocodylus niloticus',
+ 50: 'American alligator, Alligator mississipiensis',
+ 51: 'triceratops',
+ 52: 'thunder snake, worm snake, Carphophis amoenus',
+ 53: 'ringneck snake, ring-necked snake, ring snake',
+ 54: 'hognose snake, puff adder, sand viper',
+ 55: 'green snake, grass snake',
+ 56: 'king snake, kingsnake',
+ 57: 'garter snake, grass snake',
+ 58: 'water snake',
+ 59: 'vine snake',
+ 60: 'night snake, Hypsiglena torquata',
+ 61: 'boa constrictor, Constrictor constrictor',
+ 62: 'rock python, rock snake, Python sebae',
+ 63: 'Indian cobra, Naja naja',
+ 64: 'green mamba',
+ 65: 'sea snake',
+ 66: 'horned viper, cerastes, sand viper, horned asp, Cerastes cornutus',
+ 67: 'diamondback, diamondback rattlesnake, Crotalus adamanteus',
+ 68: 'sidewinder, horned rattlesnake, Crotalus cerastes',
+ 69: 'trilobite',
+ 70: 'harvestman, daddy longlegs, Phalangium opilio',
+ 71: 'scorpion',
+ 72: 'black and gold garden spider, Argiope aurantia',
+ 73: 'barn spider, Araneus cavaticus',
+ 74: 'garden spider, Aranea diademata',
+ 75: 'black widow, Latrodectus mactans',
+ 76: 'tarantula',
+ 77: 'wolf spider, hunting spider',
+ 78: 'tick',
+ 79: 'centipede',
+ 80: 'black grouse',
+ 81: 'ptarmigan',
+ 82: 'ruffed grouse, partridge, Bonasa umbellus',
+ 83: 'prairie chicken, prairie grouse, prairie fowl',
+ 84: 'peacock',
+ 85: 'quail',
+ 86: 'partridge',
+ 87: 'African grey, African gray, Psittacus erithacus',
+ 88: 'macaw',
+ 89: 'sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita',
+ 90: 'lorikeet',
+ 91: 'coucal',
+ 92: 'bee eater',
+ 93: 'hornbill',
+ 94: 'hummingbird',
+ 95: 'jacamar',
+ 96: 'toucan',
+ 97: 'drake',
+ 98: 'red-breasted merganser, Mergus serrator',
+ 99: 'goose',
+ 100: 'black swan, Cygnus atratus',
+ 101: 'tusker',
+ 102: 'echidna, spiny anteater, anteater',
+ 103: 'platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus',
+ 104: 'wallaby, brush kangaroo',
+ 105: 'koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus',
+ 106: 'wombat',
+ 107: 'jellyfish',
+ 108: 'sea anemone, anemone',
+ 109: 'brain coral',
+ 110: 'flatworm, platyhelminth',
+ 111: 'nematode, nematode worm, roundworm',
+ 112: 'conch',
+ 113: 'snail',
+ 114: 'slug',
+ 115: 'sea slug, nudibranch',
+ 116: 'chiton, coat-of-mail shell, sea cradle, polyplacophore',
+ 117: 'chambered nautilus, pearly nautilus, nautilus',
+ 118: 'Dungeness crab, Cancer magister',
+ 119: 'rock crab, Cancer irroratus',
+ 120: 'fiddler crab',
+ 121: 'king crab, Alaska crab, Alaskan king crab, Alaska king crab, Paralithodes camtschatica',
+ 122: 'American lobster, Northern lobster, Maine lobster, Homarus americanus',
+ 123: 'spiny lobster, langouste, rock lobster, crawfish, crayfish, sea crawfish',
+ 124: 'crayfish, crawfish, crawdad, crawdaddy',
+ 125: 'hermit crab',
+ 126: 'isopod',
+ 127: 'white stork, Ciconia ciconia',
+ 128: 'black stork, Ciconia nigra',
+ 129: 'spoonbill',
+ 130: 'flamingo',
+ 131: 'little blue heron, Egretta caerulea',
+ 132: 'American egret, great white heron, Egretta albus',
+ 133: 'bittern',
+ 134: 'crane',
+ 135: 'limpkin, Aramus pictus',
+ 136: 'European gallinule, Porphyrio porphyrio',
+ 137: 'American coot, marsh hen, mud hen, water hen, Fulica americana',
+ 138: 'bustard',
+ 139: 'ruddy turnstone, Arenaria interpres',
+ 140: 'red-backed sandpiper, dunlin, Erolia alpina',
+ 141: 'redshank, Tringa totanus',
+ 142: 'dowitcher',
+ 143: 'oystercatcher, oyster catcher',
+ 144: 'pelican',
+ 145: 'king penguin, Aptenodytes patagonica',
+ 146: 'albatross, mollymawk',
+ 147: 'grey whale, gray whale, devilfish, Eschrichtius gibbosus, Eschrichtius robustus',
+ 148: 'killer whale, killer, orca, grampus, sea wolf, Orcinus orca',
+ 149: 'dugong, Dugong dugon',
+ 150: 'sea lion',
+ 151: 'Chihuahua',
+ 152: 'Japanese spaniel',
+ 153: 'Maltese dog, Maltese terrier, Maltese',
+ 154: 'Pekinese, Pekingese, Peke',
+ 155: 'Shih-Tzu',
+ 156: 'Blenheim spaniel',
+ 157: 'papillon',
+ 158: 'toy terrier',
+ 159: 'Rhodesian ridgeback',
+ 160: 'Afghan hound, Afghan',
+ 161: 'basset, basset hound',
+ 162: 'beagle',
+ 163: 'bloodhound, sleuthhound',
+ 164: 'bluetick',
+ 165: 'black-and-tan coonhound',
+ 166: 'Walker hound, Walker foxhound',
+ 167: 'English foxhound',
+ 168: 'redbone',
+ 169: 'borzoi, Russian wolfhound',
+ 170: 'Irish wolfhound',
+ 171: 'Italian greyhound',
+ 172: 'whippet',
+ 173: 'Ibizan hound, Ibizan Podenco',
+ 174: 'Norwegian elkhound, elkhound',
+ 175: 'otterhound, otter hound',
+ 176: 'Saluki, gazelle hound',
+ 177: 'Scottish deerhound, deerhound',
+ 178: 'Weimaraner',
+ 179: 'Staffordshire bullterrier, Staffordshire bull terrier',
+ 180: 'American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier',
+ 181: 'Bedlington terrier',
+ 182: 'Border terrier',
+ 183: 'Kerry blue terrier',
+ 184: 'Irish terrier',
+ 185: 'Norfolk terrier',
+ 186: 'Norwich terrier',
+ 187: 'Yorkshire terrier',
+ 188: 'wire-haired fox terrier',
+ 189: 'Lakeland terrier',
+ 190: 'Sealyham terrier, Sealyham',
+ 191: 'Airedale, Airedale terrier',
+ 192: 'cairn, cairn terrier',
+ 193: 'Australian terrier',
+ 194: 'Dandie Dinmont, Dandie Dinmont terrier',
+ 195: 'Boston bull, Boston terrier',
+ 196: 'miniature schnauzer',
+ 197: 'giant schnauzer',
+ 198: 'standard schnauzer',
+ 199: 'Scotch terrier, Scottish terrier, Scottie',
+ 200: 'Tibetan terrier, chrysanthemum dog',
+ 201: 'silky terrier, Sydney silky',
+ 202: 'soft-coated wheaten terrier',
+ 203: 'West Highland white terrier',
+ 204: 'Lhasa, Lhasa apso',
+ 205: 'flat-coated retriever',
+ 206: 'curly-coated retriever',
+ 207: 'golden retriever',
+ 208: 'Labrador retriever',
+ 209: 'Chesapeake Bay retriever',
+ 210: 'German short-haired pointer',
+ 211: 'vizsla, Hungarian pointer',
+ 212: 'English setter',
+ 213: 'Irish setter, red setter',
+ 214: 'Gordon setter',
+ 215: 'Brittany spaniel',
+ 216: 'clumber, clumber spaniel',
+ 217: 'English springer, English springer spaniel',
+ 218: 'Welsh springer spaniel',
+ 219: 'cocker spaniel, English cocker spaniel, cocker',
+ 220: 'Sussex spaniel',
+ 221: 'Irish water spaniel',
+ 222: 'kuvasz',
+ 223: 'schipperke',
+ 224: 'groenendael',
+ 225: 'malinois',
+ 226: 'briard',
+ 227: 'kelpie',
+ 228: 'komondor',
+ 229: 'Old English sheepdog, bobtail',
+ 230: 'Shetland sheepdog, Shetland sheep dog, Shetland',
+ 231: 'collie',
+ 232: 'Border collie',
+ 233: 'Bouvier des Flandres, Bouviers des Flandres',
+ 234: 'Rottweiler',
+ 235: 'German shepherd, German shepherd dog, German police dog, alsatian',
+ 236: 'Doberman, Doberman pinscher',
+ 237: 'miniature pinscher',
+ 238: 'Greater Swiss Mountain dog',
+ 239: 'Bernese mountain dog',
+ 240: 'Appenzeller',
+ 241: 'EntleBucher',
+ 242: 'boxer',
+ 243: 'bull mastiff',
+ 244: 'Tibetan mastiff',
+ 245: 'French bulldog',
+ 246: 'Great Dane',
+ 247: 'Saint Bernard, St Bernard',
+ 248: 'Eskimo dog, husky',
+ 249: 'malamute, malemute, Alaskan malamute',
+ 250: 'Siberian husky',
+ 251: 'dalmatian, coach dog, carriage dog',
+ 252: 'affenpinscher, monkey pinscher, monkey dog',
+ 253: 'basenji',
+ 254: 'pug, pug-dog',
+ 255: 'Leonberg',
+ 256: 'Newfoundland, Newfoundland dog',
+ 257: 'Great Pyrenees',
+ 258: 'Samoyed, Samoyede',
+ 259: 'Pomeranian',
+ 260: 'chow, chow chow',
+ 261: 'keeshond',
+ 262: 'Brabancon griffon',
+ 263: 'Pembroke, Pembroke Welsh corgi',
+ 264: 'Cardigan, Cardigan Welsh corgi',
+ 265: 'toy poodle',
+ 266: 'miniature poodle',
+ 267: 'standard poodle',
+ 268: 'Mexican hairless',
+ 269: 'timber wolf, grey wolf, gray wolf, Canis lupus',
+ 270: 'white wolf, Arctic wolf, Canis lupus tundrarum',
+ 271: 'red wolf, maned wolf, Canis rufus, Canis niger',
+ 272: 'coyote, prairie wolf, brush wolf, Canis latrans',
+ 273: 'dingo, warrigal, warragal, Canis dingo',
+ 274: 'dhole, Cuon alpinus',
+ 275: 'African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus',
+ 276: 'hyena, hyaena',
+ 277: 'red fox, Vulpes vulpes',
+ 278: 'kit fox, Vulpes macrotis',
+ 279: 'Arctic fox, white fox, Alopex lagopus',
+ 280: 'grey fox, gray fox, Urocyon cinereoargenteus',
+ 281: 'tabby, tabby cat',
+ 282: 'tiger cat',
+ 283: 'Persian cat',
+ 284: 'Siamese cat, Siamese',
+ 285: 'Egyptian cat',
+ 286: 'cougar, puma, catamount, mountain lion, painter, panther, Felis concolor',
+ 287: 'lynx, catamount',
+ 288: 'leopard, Panthera pardus',
+ 289: 'snow leopard, ounce, Panthera uncia',
+ 290: 'jaguar, panther, Panthera onca, Felis onca',
+ 291: 'lion, king of beasts, Panthera leo',
+ 292: 'tiger, Panthera tigris',
+ 293: 'cheetah, chetah, Acinonyx jubatus',
+ 294: 'brown bear, bruin, Ursus arctos',
+ 295: 'American black bear, black bear, Ursus americanus, Euarctos americanus',
+ 296: 'ice bear, polar bear, Ursus Maritimus, Thalarctos maritimus',
+ 297: 'sloth bear, Melursus ursinus, Ursus ursinus',
+ 298: 'mongoose',
+ 299: 'meerkat, mierkat',
+ 300: 'tiger beetle',
+ 301: 'ladybug, ladybeetle, lady beetle, ladybird, ladybird beetle',
+ 302: 'ground beetle, carabid beetle',
+ 303: 'long-horned beetle, longicorn, longicorn beetle',
+ 304: 'leaf beetle, chrysomelid',
+ 305: 'dung beetle',
+ 306: 'rhinoceros beetle',
+ 307: 'weevil',
+ 308: 'fly',
+ 309: 'bee',
+ 310: 'ant, emmet, pismire',
+ 311: 'grasshopper, hopper',
+ 312: 'cricket',
+ 313: 'walking stick, walkingstick, stick insect',
+ 314: 'cockroach, roach',
+ 315: 'mantis, mantid',
+ 316: 'cicada, cicala',
+ 317: 'leafhopper',
+ 318: 'lacewing, lacewing fly',
+ 319: "dragonfly, darning needle, devil's darning needle, sewing needle, snake feeder, snake doctor, mosquito hawk, skeeter hawk",
+ 320: 'damselfly',
+ 321: 'admiral',
+ 322: 'ringlet, ringlet butterfly',
+ 323: 'monarch, monarch butterfly, milkweed butterfly, Danaus plexippus',
+ 324: 'cabbage butterfly',
+ 325: 'sulphur butterfly, sulfur butterfly',
+ 326: 'lycaenid, lycaenid butterfly',
+ 327: 'starfish, sea star',
+ 328: 'sea urchin',
+ 329: 'sea cucumber, holothurian',
+ 330: 'wood rabbit, cottontail, cottontail rabbit',
+ 331: 'hare',
+ 332: 'Angora, Angora rabbit',
+ 333: 'hamster',
+ 334: 'porcupine, hedgehog',
+ 335: 'fox squirrel, eastern fox squirrel, Sciurus niger',
+ 336: 'marmot',
+ 337: 'beaver',
+ 338: 'guinea pig, Cavia cobaya',
+ 339: 'sorrel',
+ 340: 'zebra',
+ 341: 'hog, pig, grunter, squealer, Sus scrofa',
+ 342: 'wild boar, boar, Sus scrofa',
+ 343: 'warthog',
+ 344: 'hippopotamus, hippo, river horse, Hippopotamus amphibius',
+ 345: 'ox',
+ 346: 'water buffalo, water ox, Asiatic buffalo, Bubalus bubalis',
+ 347: 'bison',
+ 348: 'ram, tup',
+ 349: 'bighorn, bighorn sheep, cimarron, Rocky Mountain bighorn, Rocky Mountain sheep, Ovis canadensis',
+ 350: 'ibex, Capra ibex',
+ 351: 'hartebeest',
+ 352: 'impala, Aepyceros melampus',
+ 353: 'gazelle',
+ 354: 'Arabian camel, dromedary, Camelus dromedarius',
+ 355: 'llama',
+ 356: 'weasel',
+ 357: 'mink',
+ 358: 'polecat, fitch, foulmart, foumart, Mustela putorius',
+ 359: 'black-footed ferret, ferret, Mustela nigripes',
+ 360: 'otter',
+ 361: 'skunk, polecat, wood pussy',
+ 362: 'badger',
+ 363: 'armadillo',
+ 364: 'three-toed sloth, ai, Bradypus tridactylus',
+ 365: 'orangutan, orang, orangutang, Pongo pygmaeus',
+ 366: 'gorilla, Gorilla gorilla',
+ 367: 'chimpanzee, chimp, Pan troglodytes',
+ 368: 'gibbon, Hylobates lar',
+ 369: 'siamang, Hylobates syndactylus, Symphalangus syndactylus',
+ 370: 'guenon, guenon monkey',
+ 371: 'patas, hussar monkey, Erythrocebus patas',
+ 372: 'baboon',
+ 373: 'macaque',
+ 374: 'langur',
+ 375: 'colobus, colobus monkey',
+ 376: 'proboscis monkey, Nasalis larvatus',
+ 377: 'marmoset',
+ 378: 'capuchin, ringtail, Cebus capucinus',
+ 379: 'howler monkey, howler',
+ 380: 'titi, titi monkey',
+ 381: 'spider monkey, Ateles geoffroyi',
+ 382: 'squirrel monkey, Saimiri sciureus',
+ 383: 'Madagascar cat, ring-tailed lemur, Lemur catta',
+ 384: 'indri, indris, Indri indri, Indri brevicaudatus',
+ 385: 'Indian elephant, Elephas maximus',
+ 386: 'African elephant, Loxodonta africana',
+ 387: 'lesser panda, red panda, panda, bear cat, cat bear, Ailurus fulgens',
+ 388: 'giant panda, panda, panda bear, coon bear, Ailuropoda melanoleuca',
+ 389: 'barracouta, snoek',
+ 390: 'eel',
+ 391: 'coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch',
+ 392: 'rock beauty, Holocanthus tricolor',
+ 393: 'anemone fish',
+ 394: 'sturgeon',
+ 395: 'gar, garfish, garpike, billfish, Lepisosteus osseus',
+ 396: 'lionfish',
+ 397: 'puffer, pufferfish, blowfish, globefish',
+ 398: 'abacus',
+ 399: 'abaya',
+ 400: "academic gown, academic robe, judge's robe",
+ 401: 'accordion, piano accordion, squeeze box',
+ 402: 'acoustic guitar',
+ 403: 'aircraft carrier, carrier, flattop, attack aircraft carrier',
+ 404: 'airliner',
+ 405: 'airship, dirigible',
+ 406: 'altar',
+ 407: 'ambulance',
+ 408: 'amphibian, amphibious vehicle',
+ 409: 'analog clock',
+ 410: 'apiary, bee house',
+ 411: 'apron',
+ 412: 'ashcan, trash can, garbage can, wastebin, ash bin, ash-bin, ashbin, dustbin, trash barrel, trash bin',
+ 413: 'assault rifle, assault gun',
+ 414: 'backpack, back pack, knapsack, packsack, rucksack, haversack',
+ 415: 'bakery, bakeshop, bakehouse',
+ 416: 'balance beam, beam',
+ 417: 'balloon',
+ 418: 'ballpoint, ballpoint pen, ballpen, Biro',
+ 419: 'Band Aid',
+ 420: 'banjo',
+ 421: 'bannister, banister, balustrade, balusters, handrail',
+ 422: 'barbell',
+ 423: 'barber chair',
+ 424: 'barbershop',
+ 425: 'barn',
+ 426: 'barometer',
+ 427: 'barrel, cask',
+ 428: 'barrow, garden cart, lawn cart, wheelbarrow',
+ 429: 'baseball',
+ 430: 'basketball',
+ 431: 'bassinet',
+ 432: 'bassoon',
+ 433: 'bathing cap, swimming cap',
+ 434: 'bath towel',
+ 435: 'bathtub, bathing tub, bath, tub',
+ 436: 'beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon',
+ 437: 'beacon, lighthouse, beacon light, pharos',
+ 438: 'beaker',
+ 439: 'bearskin, busby, shako',
+ 440: 'beer bottle',
+ 441: 'beer glass',
+ 442: 'bell cote, bell cot',
+ 443: 'bib',
+ 444: 'bicycle-built-for-two, tandem bicycle, tandem',
+ 445: 'bikini, two-piece',
+ 446: 'binder, ring-binder',
+ 447: 'binoculars, field glasses, opera glasses',
+ 448: 'birdhouse',
+ 449: 'boathouse',
+ 450: 'bobsled, bobsleigh, bob',
+ 451: 'bolo tie, bolo, bola tie, bola',
+ 452: 'bonnet, poke bonnet',
+ 453: 'bookcase',
+ 454: 'bookshop, bookstore, bookstall',
+ 455: 'bottlecap',
+ 456: 'bow',
+ 457: 'bow tie, bow-tie, bowtie',
+ 458: 'brass, memorial tablet, plaque',
+ 459: 'brassiere, bra, bandeau',
+ 460: 'breakwater, groin, groyne, mole, bulwark, seawall, jetty',
+ 461: 'breastplate, aegis, egis',
+ 462: 'broom',
+ 463: 'bucket, pail',
+ 464: 'buckle',
+ 465: 'bulletproof vest',
+ 466: 'bullet train, bullet',
+ 467: 'butcher shop, meat market',
+ 468: 'cab, hack, taxi, taxicab',
+ 469: 'caldron, cauldron',
+ 470: 'candle, taper, wax light',
+ 471: 'cannon',
+ 472: 'canoe',
+ 473: 'can opener, tin opener',
+ 474: 'cardigan',
+ 475: 'car mirror',
+ 476: 'carousel, carrousel, merry-go-round, roundabout, whirligig',
+ 477: "carpenter's kit, tool kit",
+ 478: 'carton',
+ 479: 'car wheel',
+ 480: 'cash machine, cash dispenser, automated teller machine, automatic teller machine, automated teller, automatic teller, ATM',
+ 481: 'cassette',
+ 482: 'cassette player',
+ 483: 'castle',
+ 484: 'catamaran',
+ 485: 'CD player',
+ 486: 'cello, violoncello',
+ 487: 'cellular telephone, cellular phone, cellphone, cell, mobile phone',
+ 488: 'chain',
+ 489: 'chainlink fence',
+ 490: 'chain mail, ring mail, mail, chain armor, chain armour, ring armor, ring armour',
+ 491: 'chain saw, chainsaw',
+ 492: 'chest',
+ 493: 'chiffonier, commode',
+ 494: 'chime, bell, gong',
+ 495: 'china cabinet, china closet',
+ 496: 'Christmas stocking',
+ 497: 'church, church building',
+ 498: 'cinema, movie theater, movie theatre, movie house, picture palace',
+ 499: 'cleaver, meat cleaver, chopper',
+ 500: 'cliff dwelling',
+ 501: 'cloak',
+ 502: 'clog, geta, patten, sabot',
+ 503: 'cocktail shaker',
+ 504: 'coffee mug',
+ 505: 'coffeepot',
+ 506: 'coil, spiral, volute, whorl, helix',
+ 507: 'combination lock',
+ 508: 'computer keyboard, keypad',
+ 509: 'confectionery, confectionary, candy store',
+ 510: 'container ship, containership, container vessel',
+ 511: 'convertible',
+ 512: 'corkscrew, bottle screw',
+ 513: 'cornet, horn, trumpet, trump',
+ 514: 'cowboy boot',
+ 515: 'cowboy hat, ten-gallon hat',
+ 516: 'cradle',
+ 517: 'crane',
+ 518: 'crash helmet',
+ 519: 'crate',
+ 520: 'crib, cot',
+ 521: 'Crock Pot',
+ 522: 'croquet ball',
+ 523: 'crutch',
+ 524: 'cuirass',
+ 525: 'dam, dike, dyke',
+ 526: 'desk',
+ 527: 'desktop computer',
+ 528: 'dial telephone, dial phone',
+ 529: 'diaper, nappy, napkin',
+ 530: 'digital clock',
+ 531: 'digital watch',
+ 532: 'dining table, board',
+ 533: 'dishrag, dishcloth',
+ 534: 'dishwasher, dish washer, dishwashing machine',
+ 535: 'disk brake, disc brake',
+ 536: 'dock, dockage, docking facility',
+ 537: 'dogsled, dog sled, dog sleigh',
+ 538: 'dome',
+ 539: 'doormat, welcome mat',
+ 540: 'drilling platform, offshore rig',
+ 541: 'drum, membranophone, tympan',
+ 542: 'drumstick',
+ 543: 'dumbbell',
+ 544: 'Dutch oven',
+ 545: 'electric fan, blower',
+ 546: 'electric guitar',
+ 547: 'electric locomotive',
+ 548: 'entertainment center',
+ 549: 'envelope',
+ 550: 'espresso maker',
+ 551: 'face powder',
+ 552: 'feather boa, boa',
+ 553: 'file, file cabinet, filing cabinet',
+ 554: 'fireboat',
+ 555: 'fire engine, fire truck',
+ 556: 'fire screen, fireguard',
+ 557: 'flagpole, flagstaff',
+ 558: 'flute, transverse flute',
+ 559: 'folding chair',
+ 560: 'football helmet',
+ 561: 'forklift',
+ 562: 'fountain',
+ 563: 'fountain pen',
+ 564: 'four-poster',
+ 565: 'freight car',
+ 566: 'French horn, horn',
+ 567: 'frying pan, frypan, skillet',
+ 568: 'fur coat',
+ 569: 'garbage truck, dustcart',
+ 570: 'gasmask, respirator, gas helmet',
+ 571: 'gas pump, gasoline pump, petrol pump, island dispenser',
+ 572: 'goblet',
+ 573: 'go-kart',
+ 574: 'golf ball',
+ 575: 'golfcart, golf cart',
+ 576: 'gondola',
+ 577: 'gong, tam-tam',
+ 578: 'gown',
+ 579: 'grand piano, grand',
+ 580: 'greenhouse, nursery, glasshouse',
+ 581: 'grille, radiator grille',
+ 582: 'grocery store, grocery, food market, market',
+ 583: 'guillotine',
+ 584: 'hair slide',
+ 585: 'hair spray',
+ 586: 'half track',
+ 587: 'hammer',
+ 588: 'hamper',
+ 589: 'hand blower, blow dryer, blow drier, hair dryer, hair drier',
+ 590: 'hand-held computer, hand-held microcomputer',
+ 591: 'handkerchief, hankie, hanky, hankey',
+ 592: 'hard disc, hard disk, fixed disk',
+ 593: 'harmonica, mouth organ, harp, mouth harp',
+ 594: 'harp',
+ 595: 'harvester, reaper',
+ 596: 'hatchet',
+ 597: 'holster',
+ 598: 'home theater, home theatre',
+ 599: 'honeycomb',
+ 600: 'hook, claw',
+ 601: 'hoopskirt, crinoline',
+ 602: 'horizontal bar, high bar',
+ 603: 'horse cart, horse-cart',
+ 604: 'hourglass',
+ 605: 'iPod',
+ 606: 'iron, smoothing iron',
+ 607: "jack-o'-lantern",
+ 608: 'jean, blue jean, denim',
+ 609: 'jeep, landrover',
+ 610: 'jersey, T-shirt, tee shirt',
+ 611: 'jigsaw puzzle',
+ 612: 'jinrikisha, ricksha, rickshaw',
+ 613: 'joystick',
+ 614: 'kimono',
+ 615: 'knee pad',
+ 616: 'knot',
+ 617: 'lab coat, laboratory coat',
+ 618: 'ladle',
+ 619: 'lampshade, lamp shade',
+ 620: 'laptop, laptop computer',
+ 621: 'lawn mower, mower',
+ 622: 'lens cap, lens cover',
+ 623: 'letter opener, paper knife, paperknife',
+ 624: 'library',
+ 625: 'lifeboat',
+ 626: 'lighter, light, igniter, ignitor',
+ 627: 'limousine, limo',
+ 628: 'liner, ocean liner',
+ 629: 'lipstick, lip rouge',
+ 630: 'Loafer',
+ 631: 'lotion',
+ 632: 'loudspeaker, speaker, speaker unit, loudspeaker system, speaker system',
+ 633: "loupe, jeweler's loupe",
+ 634: 'lumbermill, sawmill',
+ 635: 'magnetic compass',
+ 636: 'mailbag, postbag',
+ 637: 'mailbox, letter box',
+ 638: 'maillot',
+ 639: 'maillot, tank suit',
+ 640: 'manhole cover',
+ 641: 'maraca',
+ 642: 'marimba, xylophone',
+ 643: 'mask',
+ 644: 'matchstick',
+ 645: 'maypole',
+ 646: 'maze, labyrinth',
+ 647: 'measuring cup',
+ 648: 'medicine chest, medicine cabinet',
+ 649: 'megalith, megalithic structure',
+ 650: 'microphone, mike',
+ 651: 'microwave, microwave oven',
+ 652: 'military uniform',
+ 653: 'milk can',
+ 654: 'minibus',
+ 655: 'miniskirt, mini',
+ 656: 'minivan',
+ 657: 'missile',
+ 658: 'mitten',
+ 659: 'mixing bowl',
+ 660: 'mobile home, manufactured home',
+ 661: 'Model T',
+ 662: 'modem',
+ 663: 'monastery',
+ 664: 'monitor',
+ 665: 'moped',
+ 666: 'mortar',
+ 667: 'mortarboard',
+ 668: 'mosque',
+ 669: 'mosquito net',
+ 670: 'motor scooter, scooter',
+ 671: 'mountain bike, all-terrain bike, off-roader',
+ 672: 'mountain tent',
+ 673: 'mouse, computer mouse',
+ 674: 'mousetrap',
+ 675: 'moving van',
+ 676: 'muzzle',
+ 677: 'nail',
+ 678: 'neck brace',
+ 679: 'necklace',
+ 680: 'nipple',
+ 681: 'notebook, notebook computer',
+ 682: 'obelisk',
+ 683: 'oboe, hautboy, hautbois',
+ 684: 'ocarina, sweet potato',
+ 685: 'odometer, hodometer, mileometer, milometer',
+ 686: 'oil filter',
+ 687: 'organ, pipe organ',
+ 688: 'oscilloscope, scope, cathode-ray oscilloscope, CRO',
+ 689: 'overskirt',
+ 690: 'oxcart',
+ 691: 'oxygen mask',
+ 692: 'packet',
+ 693: 'paddle, boat paddle',
+ 694: 'paddlewheel, paddle wheel',
+ 695: 'padlock',
+ 696: 'paintbrush',
+ 697: "pajama, pyjama, pj's, jammies",
+ 698: 'palace',
+ 699: 'panpipe, pandean pipe, syrinx',
+ 700: 'paper towel',
+ 701: 'parachute, chute',
+ 702: 'parallel bars, bars',
+ 703: 'park bench',
+ 704: 'parking meter',
+ 705: 'passenger car, coach, carriage',
+ 706: 'patio, terrace',
+ 707: 'pay-phone, pay-station',
+ 708: 'pedestal, plinth, footstall',
+ 709: 'pencil box, pencil case',
+ 710: 'pencil sharpener',
+ 711: 'perfume, essence',
+ 712: 'Petri dish',
+ 713: 'photocopier',
+ 714: 'pick, plectrum, plectron',
+ 715: 'pickelhaube',
+ 716: 'picket fence, paling',
+ 717: 'pickup, pickup truck',
+ 718: 'pier',
+ 719: 'piggy bank, penny bank',
+ 720: 'pill bottle',
+ 721: 'pillow',
+ 722: 'ping-pong ball',
+ 723: 'pinwheel',
+ 724: 'pirate, pirate ship',
+ 725: 'pitcher, ewer',
+ 726: "plane, carpenter's plane, woodworking plane",
+ 727: 'planetarium',
+ 728: 'plastic bag',
+ 729: 'plate rack',
+ 730: 'plow, plough',
+ 731: "plunger, plumber's helper",
+ 732: 'Polaroid camera, Polaroid Land camera',
+ 733: 'pole',
+ 734: 'police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria',
+ 735: 'poncho',
+ 736: 'pool table, billiard table, snooker table',
+ 737: 'pop bottle, soda bottle',
+ 738: 'pot, flowerpot',
+ 739: "potter's wheel",
+ 740: 'power drill',
+ 741: 'prayer rug, prayer mat',
+ 742: 'printer',
+ 743: 'prison, prison house',
+ 744: 'projectile, missile',
+ 745: 'projector',
+ 746: 'puck, hockey puck',
+ 747: 'punching bag, punch bag, punching ball, punchball',
+ 748: 'purse',
+ 749: 'quill, quill pen',
+ 750: 'quilt, comforter, comfort, puff',
+ 751: 'racer, race car, racing car',
+ 752: 'racket, racquet',
+ 753: 'radiator',
+ 754: 'radio, wireless',
+ 755: 'radio telescope, radio reflector',
+ 756: 'rain barrel',
+ 757: 'recreational vehicle, RV, R.V.',
+ 758: 'reel',
+ 759: 'reflex camera',
+ 760: 'refrigerator, icebox',
+ 761: 'remote control, remote',
+ 762: 'restaurant, eating house, eating place, eatery',
+ 763: 'revolver, six-gun, six-shooter',
+ 764: 'rifle',
+ 765: 'rocking chair, rocker',
+ 766: 'rotisserie',
+ 767: 'rubber eraser, rubber, pencil eraser',
+ 768: 'rugby ball',
+ 769: 'rule, ruler',
+ 770: 'running shoe',
+ 771: 'safe',
+ 772: 'safety pin',
+ 773: 'saltshaker, salt shaker',
+ 774: 'sandal',
+ 775: 'sarong',
+ 776: 'sax, saxophone',
+ 777: 'scabbard',
+ 778: 'scale, weighing machine',
+ 779: 'school bus',
+ 780: 'schooner',
+ 781: 'scoreboard',
+ 782: 'screen, CRT screen',
+ 783: 'screw',
+ 784: 'screwdriver',
+ 785: 'seat belt, seatbelt',
+ 786: 'sewing machine',
+ 787: 'shield, buckler',
+ 788: 'shoe shop, shoe-shop, shoe store',
+ 789: 'shoji',
+ 790: 'shopping basket',
+ 791: 'shopping cart',
+ 792: 'shovel',
+ 793: 'shower cap',
+ 794: 'shower curtain',
+ 795: 'ski',
+ 796: 'ski mask',
+ 797: 'sleeping bag',
+ 798: 'slide rule, slipstick',
+ 799: 'sliding door',
+ 800: 'slot, one-armed bandit',
+ 801: 'snorkel',
+ 802: 'snowmobile',
+ 803: 'snowplow, snowplough',
+ 804: 'soap dispenser',
+ 805: 'soccer ball',
+ 806: 'sock',
+ 807: 'solar dish, solar collector, solar furnace',
+ 808: 'sombrero',
+ 809: 'soup bowl',
+ 810: 'space bar',
+ 811: 'space heater',
+ 812: 'space shuttle',
+ 813: 'spatula',
+ 814: 'speedboat',
+ 815: "spider web, spider's web",
+ 816: 'spindle',
+ 817: 'sports car, sport car',
+ 818: 'spotlight, spot',
+ 819: 'stage',
+ 820: 'steam locomotive',
+ 821: 'steel arch bridge',
+ 822: 'steel drum',
+ 823: 'stethoscope',
+ 824: 'stole',
+ 825: 'stone wall',
+ 826: 'stopwatch, stop watch',
+ 827: 'stove',
+ 828: 'strainer',
+ 829: 'streetcar, tram, tramcar, trolley, trolley car',
+ 830: 'stretcher',
+ 831: 'studio couch, day bed',
+ 832: 'stupa, tope',
+ 833: 'submarine, pigboat, sub, U-boat',
+ 834: 'suit, suit of clothes',
+ 835: 'sundial',
+ 836: 'sunglass',
+ 837: 'sunglasses, dark glasses, shades',
+ 838: 'sunscreen, sunblock, sun blocker',
+ 839: 'suspension bridge',
+ 840: 'swab, swob, mop',
+ 841: 'sweatshirt',
+ 842: 'swimming trunks, bathing trunks',
+ 843: 'swing',
+ 844: 'switch, electric switch, electrical switch',
+ 845: 'syringe',
+ 846: 'table lamp',
+ 847: 'tank, army tank, armored combat vehicle, armoured combat vehicle',
+ 848: 'tape player',
+ 849: 'teapot',
+ 850: 'teddy, teddy bear',
+ 851: 'television, television system',
+ 852: 'tennis ball',
+ 853: 'thatch, thatched roof',
+ 854: 'theater curtain, theatre curtain',
+ 855: 'thimble',
+ 856: 'thresher, thrasher, threshing machine',
+ 857: 'throne',
+ 858: 'tile roof',
+ 859: 'toaster',
+ 860: 'tobacco shop, tobacconist shop, tobacconist',
+ 861: 'toilet seat',
+ 862: 'torch',
+ 863: 'totem pole',
+ 864: 'tow truck, tow car, wrecker',
+ 865: 'toyshop',
+ 866: 'tractor',
+ 867: 'trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi',
+ 868: 'tray',
+ 869: 'trench coat',
+ 870: 'tricycle, trike, velocipede',
+ 871: 'trimaran',
+ 872: 'tripod',
+ 873: 'triumphal arch',
+ 874: 'trolleybus, trolley coach, trackless trolley',
+ 875: 'trombone',
+ 876: 'tub, vat',
+ 877: 'turnstile',
+ 878: 'typewriter keyboard',
+ 879: 'umbrella',
+ 880: 'unicycle, monocycle',
+ 881: 'upright, upright piano',
+ 882: 'vacuum, vacuum cleaner',
+ 883: 'vase',
+ 884: 'vault',
+ 885: 'velvet',
+ 886: 'vending machine',
+ 887: 'vestment',
+ 888: 'viaduct',
+ 889: 'violin, fiddle',
+ 890: 'volleyball',
+ 891: 'waffle iron',
+ 892: 'wall clock',
+ 893: 'wallet, billfold, notecase, pocketbook',
+ 894: 'wardrobe, closet, press',
+ 895: 'warplane, military plane',
+ 896: 'washbasin, handbasin, washbowl, lavabo, wash-hand basin',
+ 897: 'washer, automatic washer, washing machine',
+ 898: 'water bottle',
+ 899: 'water jug',
+ 900: 'water tower',
+ 901: 'whiskey jug',
+ 902: 'whistle',
+ 903: 'wig',
+ 904: 'window screen',
+ 905: 'window shade',
+ 906: 'Windsor tie',
+ 907: 'wine bottle',
+ 908: 'wing',
+ 909: 'wok',
+ 910: 'wooden spoon',
+ 911: 'wool, woolen, woollen',
+ 912: 'worm fence, snake fence, snake-rail fence, Virginia fence',
+ 913: 'wreck',
+ 914: 'yawl',
+ 915: 'yurt',
+ 916: 'web site, website, internet site, site',
+ 917: 'comic book',
+ 918: 'crossword puzzle, crossword',
+ 919: 'street sign',
+ 920: 'traffic light, traffic signal, stoplight',
+ 921: 'book jacket, dust cover, dust jacket, dust wrapper',
+ 922: 'menu',
+ 923: 'plate',
+ 924: 'guacamole',
+ 925: 'consomme',
+ 926: 'hot pot, hotpot',
+ 927: 'trifle',
+ 928: 'ice cream, icecream',
+ 929: 'ice lolly, lolly, lollipop, popsicle',
+ 930: 'French loaf',
+ 931: 'bagel, beigel',
+ 932: 'pretzel',
+ 933: 'cheeseburger',
+ 934: 'hotdog, hot dog, red hot',
+ 935: 'mashed potato',
+ 936: 'head cabbage',
+ 937: 'broccoli',
+ 938: 'cauliflower',
+ 939: 'zucchini, courgette',
+ 940: 'spaghetti squash',
+ 941: 'acorn squash',
+ 942: 'butternut squash',
+ 943: 'cucumber, cuke',
+ 944: 'artichoke, globe artichoke',
+ 945: 'bell pepper',
+ 946: 'cardoon',
+ 947: 'mushroom',
+ 948: 'Granny Smith',
+ 949: 'strawberry',
+ 950: 'orange',
+ 951: 'lemon',
+ 952: 'fig',
+ 953: 'pineapple, ananas',
+ 954: 'banana',
+ 955: 'jackfruit, jak, jack',
+ 956: 'custard apple',
+ 957: 'pomegranate',
+ 958: 'hay',
+ 959: 'carbonara',
+ 960: 'chocolate sauce, chocolate syrup',
+ 961: 'dough',
+ 962: 'meat loaf, meatloaf',
+ 963: 'pizza, pizza pie',
+ 964: 'potpie',
+ 965: 'burrito',
+ 966: 'red wine',
+ 967: 'espresso',
+ 968: 'cup',
+ 969: 'eggnog',
+ 970: 'alp',
+ 971: 'bubble',
+ 972: 'cliff, drop, drop-off',
+ 973: 'coral reef',
+ 974: 'geyser',
+ 975: 'lakeside, lakeshore',
+ 976: 'promontory, headland, head, foreland',
+ 977: 'sandbar, sand bar',
+ 978: 'seashore, coast, seacoast, sea-coast',
+ 979: 'valley, vale',
+ 980: 'volcano',
+ 981: 'ballplayer, baseball player',
+ 982: 'groom, bridegroom',
+ 983: 'scuba diver',
+ 984: 'rapeseed',
+ 985: 'daisy',
+ 986: "yellow lady's slipper, yellow lady-slipper, Cypripedium calceolus, Cypripedium parviflorum",
+ 987: 'corn',
+ 988: 'acorn',
+ 989: 'hip, rose hip, rosehip',
+ 990: 'buckeye, horse chestnut, conker',
+ 991: 'coral fungus',
+ 992: 'agaric',
+ 993: 'gyromitra',
+ 994: 'stinkhorn, carrion fungus',
+ 995: 'earthstar',
+ 996: 'hen-of-the-woods, hen of the woods, Polyporus frondosus, Grifola frondosa',
+ 997: 'bolete',
+ 998: 'ear, spike, capitulum',
+ 999: 'toilet tissue, toilet paper, bathroom tissue'}

modelguidedattacks/data/registry.py

@@ -0,0 +1,108 @@
+import torch
+import torchvision
+import torchvision.transforms as T
+from mmpretrain import datasets as mmdatasets
+from mmpretrain.registry import TRANSFORMS
+from mmengine.dataset import Compose
+
+from torch import nn
+from torch.utils.data import Dataset as TorchDataset
+
+# This holds dataset instantiation functions by (dataset_name) tuple keys
+DATASET_REGISTRY = {}
+DATASET_PATH = "./datasets"
+
+class MMPretrainWrapper(TorchDataset):
+    def __init__(self, mmdataset) -> None:
+        super().__init__()
+        self.mmdataset = mmdataset
+
+        test_pipeline = [
+            dict(type='LoadImageFromFile'),
+            dict(type='ResizeEdge', scale=256, edge='short'),
+            dict(type='CenterCrop', crop_size=224),
+            dict(type='PackInputs'),
+        ]
+
+        self.pipeline = self.init_pipeline(test_pipeline)
+
+    def init_pipeline(self, pipeline_cfg):
+        pipeline = Compose(
+            [TRANSFORMS.build(t) for t in pipeline_cfg])
+        return pipeline
+
+    @property
+    def classes(self):
+        return self.mmdataset.CLASSES
+    
+    def __len__(self):
+        return len(self.mmdataset)
+
+    def __getitem__(self, index):
+        sample = self.mmdataset[index]
+        sample = self.pipeline(sample)
+
+        # Our interface expects images in [0-1]
+        img = sample["inputs"].float() / 255
+
+        return img, sample["data_samples"].gt_label.item()
+
+
+def register_torchvision_dataset(dataset_name, dataset_cls, dataset_kwargs_train={}, dataset_kwargs_val={}):
+    def instantiate_dataset():
+        train_data = dataset_cls(
+            root=DATASET_PATH,
+            train=True,
+            download=True,
+            transform=T.ToTensor()
+        )
+
+        val_data = dataset_cls(
+            root=DATASET_PATH,
+            train=False,
+            download=True,
+            transform=T.ToTensor()
+        )
+
+        return train_data, val_data
+
+    DATASET_REGISTRY[dataset_name] = instantiate_dataset
+
+def register_mmpretrain_dataset(dataset_name, dataset_cls, dataset_kwargs_train={}, dataset_kwargs_val={}):
+    def instantiate_dataset():
+        train_data = dataset_cls(**dataset_kwargs_train)
+        val_data = dataset_cls(**dataset_kwargs_val)
+
+        train_data = MMPretrainWrapper(train_data)
+        val_data = MMPretrainWrapper(val_data)
+
+        return train_data, val_data
+    
+    DATASET_REGISTRY[dataset_name] = instantiate_dataset
+
+def register_default_datasets():
+    register_torchvision_dataset("cifar10", torchvision.datasets.CIFAR10)
+    register_torchvision_dataset("cifar100", torchvision.datasets.CIFAR100)
+    register_mmpretrain_dataset("imagenet", mmdatasets.ImageNet, 
+                                dataset_kwargs_train=dict(
+                                    data_root = "data/imagenet", 
+                                    data_prefix = "val", 
+                                    ann_file = "meta/val.txt"
+                                ),
+                                dataset_kwargs_val=dict(
+                                    data_root = "data/imagenet", 
+                                    data_prefix = "val", 
+                                    ann_file = "meta/val.txt"
+                                ))
+
+def get_dataset(dataset_name):
+    """
+    Returns an instance of a dataset
+
+    dataset_name: Name of desired dataset
+    """
+
+    if dataset_name not in DATASET_REGISTRY:
+        raise Exception("Requested dataset not in registry")        
+    
+    return DATASET_REGISTRY[dataset_name]()

modelguidedattacks/data/setup.py

@@ -0,0 +1,170 @@
+from typing import Any
+import os
+import torch
+import ignite.distributed as idist
+import torchvision
+import torchvision.transforms as T
+from torch.utils import data as torch_data
+
+from .classification_wrapper import TopKClassificationWrapper
+from torch.utils.data import Subset
+from modelguidedattacks.data import get_dataset
+from modelguidedattacks.cls_models.accuracy import get_correct_subset_for_models, DATASET_METADATA_DIR
+
+from tqdm import tqdm
+
+def get_gt_labels(dataset: TopKClassificationWrapper, train:bool, dataset_name:str):
+    training_str = "train" if train else "val"
+    save_name = os.path.join(DATASET_METADATA_DIR, f"{dataset_name}_labels_{training_str}.p")
+
+    if os.path.exists(save_name):
+        print ("Found labels cache")
+        return torch.load(save_name)
+
+    dataloader = torch_data.DataLoader(dataset, batch_size=128, shuffle=False, num_workers=4)
+
+    gt_labels = []
+
+    for batch in tqdm(dataloader):
+        gt_labels.extend(batch[1].tolist())
+
+    gt_labels = torch.tensor(gt_labels)
+
+    torch.save(gt_labels, save_name)
+
+    return gt_labels
+
+def class_balanced_sampling(dataset, gt_labels: torch.Tensor,
+                            correct_labels: list, total_samples=1000):
+    num_classes = len(dataset.classes)
+
+    correct_labels = torch.tensor(correct_labels)
+    correct_mask = torch.zeros((len(dataset), ), dtype=torch.bool)
+    correct_mask[correct_labels] = True
+
+    sampled_indices = 0
+
+    total_sampled_indices = 0
+    sampled_indices = [[] for i in range(num_classes)]
+
+    shuffled_inds = torch.randperm(len(dataset))
+
+    for sample_cnt, sample_i in enumerate(shuffled_inds):
+        if not correct_mask[sample_i]:
+            continue
+
+        sample_class = gt_labels[sample_i]
+        desired_samples_in_class = (total_sampled_indices // num_classes) + 1
+
+        if len(sampled_indices[sample_class]) < desired_samples_in_class:
+            sampled_indices[sample_class].append(sample_i.item())
+            total_sampled_indices += 1
+
+            if total_sampled_indices >= total_samples:
+                break
+
+    flattened_indices = []
+    for class_samples in sampled_indices:
+        flattened_indices.extend(class_samples)
+
+    return torch.tensor(flattened_indices)
+
+def sample_attack_labels(dataset, gt_labels, k, sampler):
+    """
+    dataset: Dataset we're generating attack labels for
+    gt_labels: List of gt idx for each sample in a dataset
+    k: attack size
+    sampler: ["random"]
+    """
+
+    # Sample from uniform and argsort to simulate
+    # a batched randperm
+    attack_label_uniforms = torch.rand((len(gt_labels), len(dataset.classes)))
+
+    # We don't want to sample the gt class for any samples
+    batch_inds = torch.arange(len(gt_labels))
+    attack_label_uniforms[batch_inds, gt_labels] = -1.
+
+    attack_labels = attack_label_uniforms.argsort(dim=-1, descending=True)[:, :k]
+
+    return attack_labels
+
+def setup_data(config: Any, rank):
+    """Download datasets and create dataloaders
+
+    Parameters
+    ----------
+    config: needs to contain `data_path`, `train_batch_size`, `eval_batch_size`, and `num_workers`
+    """
+
+    dataset_train, dataset_eval = get_dataset(config.dataset)
+
+    train_subset = None
+    val_subset = None
+
+    attack_labels_train = None
+    attack_labels_val = None
+
+    if rank == 0:
+        gt_labels_train = get_gt_labels(dataset_train, True, config.dataset)
+        gt_labels_val = get_gt_labels(dataset_eval, False, config.dataset)
+
+        attack_labels_train = sample_attack_labels(dataset_train, gt_labels_train, k=config.k,
+                                                   sampler=config.attack_sampling)
+        attack_labels_val = sample_attack_labels(dataset_eval, gt_labels_val, k=config.k,
+                                                 sampler=config.attack_sampling)
+    
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        correct_train_set = get_correct_subset_for_models(config.compare_models, 
+                                                        config.dataset, device, 
+                                                        train=True)
+        
+        correct_eval_set = get_correct_subset_for_models(config.compare_models, 
+                                                        config.dataset, device, 
+                                                        train=False)
+        
+        # Balanced sampling
+        train_subset = class_balanced_sampling(dataset_train, gt_labels_train,
+                                               correct_train_set)
+        
+        val_subset = class_balanced_sampling(dataset_eval, gt_labels_val,
+                                             correct_eval_set)
+        
+        if config.overfit:
+            rand_inds = torch.randperm(len(val_subset))[:16]
+            train_subset = train_subset[rand_inds]
+            val_subset = val_subset[rand_inds]
+    
+    train_subset = idist.broadcast(train_subset, safe_mode=True)
+    val_subset = idist.broadcast(val_subset, safe_mode=True)
+
+    attack_labels_train = idist.broadcast(attack_labels_train, safe_mode=True)
+    attack_labels_val = idist.broadcast(attack_labels_val, safe_mode=True)
+
+    dataset_train = TopKClassificationWrapper(dataset_train, k=config.k, 
+                                              attack_labels=attack_labels_train)
+    dataset_eval = TopKClassificationWrapper(dataset_eval, k=config.k, 
+                                             attack_labels=attack_labels_val)
+
+    dataset_train = Subset(dataset_train, train_subset)
+    dataset_eval = Subset(dataset_eval, val_subset)
+
+    # if config.overfit:
+    #     dataset_train = Subset(dataset_train, range(2))
+    #     dataset_eval = dataset_train
+    # else:
+    #     dataset_eval = Subset(dataset_eval, torch.randperm(len(dataset_eval))[:1000].tolist() )
+
+    dataloader_train = idist.auto_dataloader(
+        dataset_train,
+        batch_size=config.train_batch_size,
+        shuffle=not config.overfit,
+        num_workers=config.num_workers,
+    )
+    dataloader_eval = idist.auto_dataloader(
+        dataset_eval,
+        batch_size=config.eval_batch_size,
+        shuffle=True,
+        num_workers=config.num_workers,
+    )
+    return dataloader_train, dataloader_eval

modelguidedattacks/guides/instance_guide.py

@@ -0,0 +1,109 @@
+import torch
+from torch import nn
+from torchvision.ops import MLP
+
+from .. import losses
+
+class InstanceGuide(nn.Module):
+    def __init__(self, model: nn.Module, optimizer=torch.optim.AdamW, loss_fn=losses.CWExtensionLoss) -> None:
+        super().__init__()
+        
+        self.guided = True
+        self.model = model
+        
+
+        for p in self.model.parameters():
+            p.requires_grad_(False)
+
+        self.loss = loss_fn()
+        self.optimizer = optimizer
+
+        self.epochs = 30
+        self.mlp_iterations = 5
+        self.perturbation_iterations = 5
+
+    def surject_perturbation(self, x):
+        return x
+
+    def forward(self, x, attack_targets):
+        """
+        x: [B, channels, H, W]
+        attack_targets: [B, K]
+        """
+
+        B = x.shape[0]
+        K = attack_targets.shape[-1]
+        C = self.model.num_classes()
+
+        with torch.no_grad():
+            pred_clean, feats = self.model(x, return_features=True)
+
+        # We are assuming the clean predictions are ground truth since we make that
+        # constraint on the dataset side
+        attack_ground_truth = pred_clean.argmax(dim=-1) # [B]
+
+        mlp = MLP(self.model.head_features(), 
+                    [self.model.head_features()]*3 + [self.model.head_features()],
+                    activation_layer=nn.GELU, inplace=None).to(x.device)
+
+        x_perturbation = nn.Parameter(torch.randn(x.shape, 
+                                                 device=x.device)*1e-3)
+        
+        perturbation_optimizer = self.optimizer([x_perturbation], lr=1e-1)
+
+        mlp_optimizer = self.optimizer(mlp.parameters(), lr=1e-3)
+
+        logits_target_best = pred_clean
+        feats_target_best = feats
+
+        with torch.enable_grad():
+            for i in range(self.epochs):
+                for _ in range(self.mlp_iterations):
+                    torch.cuda.synchronize()
+
+                    feature_offset = mlp(feats)
+                    feats_target_pred = feature_offset + feats
+                    logits_target_pred = self.model.head(feats_target_pred)
+                    # logits_target_pred = pred_logits
+                    pred_classes = logits_target_pred.argsort(dim=-1, descending=True) # [B, C]
+                    attack_successful = (pred_classes[:, :K] == attack_targets).all(dim=-1) # [B]
+
+                    with torch.no_grad():
+                        logits_target_best = torch.where(
+                            attack_successful[:, None].expand(-1, C),
+                            logits_target_pred,
+                            logits_target_best
+                        )
+
+                        feats_target_best = torch.where(
+                            attack_successful[:, None].expand(-1, self.model.head_features()),
+                            feats_target_pred,
+                            feats_target_best
+                        )
+
+                    mlp_loss = self.loss(logits_pred=logits_target_pred, 
+                                         prediction_feats=feats_target_pred,
+                                         attack_targets=attack_targets, 
+                                         attack_ground_truth=attack_ground_truth, 
+                                         model=self.model)
+                    mlp_loss = mlp_loss.mean() + feature_offset.view(B, -1).norm(dim=-1, p=2)*1
+
+                    mlp_optimizer.zero_grad()
+                    mlp_loss.backward()
+                    mlp_optimizer.step()
+
+                feats_target_best = feats_target_best.detach()
+
+                for _ in range(self.perturbation_iterations):
+                    x_perturbed = x + self.surject_perturbation(x_perturbation)
+                    prediction, perturbed_feats = self.model(x_perturbed, return_features=True)
+                    pred_classes = prediction.argsort(dim=-1, descending=True) # [B, C]
+                    attack_successful = (pred_classes[:, :K] == attack_targets).all(dim=-1) # [B]
+
+                    perturbation_loss = (prediction - logits_target_best).view(B, -1).norm(dim=-1).mean()
+
+                    perturbation_optimizer.zero_grad()
+                    perturbation_loss.backward()
+                    perturbation_optimizer.step()
+            
+        return prediction

modelguidedattacks/guides/unguided.py

@@ -0,0 +1,314 @@
+import torch
+from torch import nn
+from .. import losses
+import ignite.distributed as idist
+import torch_optimizer 
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+from torch.nn import functional as F
+import os
+
+import shutil
+from modelguidedattacks.cls_models.registry import MMPretrainVisualTransformerWrapper
+from modelguidedattacks.data.imagenet_metadata import imgnet_idx_to_name
+
+class Unguided(nn.Module):
+    def __init__(self, model: nn.Module, config, optimizer=torch.optim.AdamW, seed=0, iterations=1000,
+                 loss_fn=losses.CVXProjLoss, lr=1e-3, 
+                 binary_search_steps=1, topk_loss_coef_upper=10.,
+                 topk_loss_coef_lower=0.) -> None:
+        super().__init__()
+
+        self.guided = False
+        self.model = model
+        self.seed = seed
+        self.iterations = iterations
+        self.loss = loss_fn()
+        self.optimizer = optimizer
+        self.lr = lr
+
+        self.binary_search_steps = binary_search_steps
+        self.topk_loss_coef_upper = topk_loss_coef_upper
+        self.topk_loss_coef_lower = topk_loss_coef_lower
+        self.config = config
+
+    def surject_perturbation(self, x, max_norm=5.):
+        x_shape = x.shape
+
+        x = x.flatten(1)
+        x_norm = x.norm(dim=-1)
+        x_unit = x / x_norm[:, None]
+
+        x_norm_outside = x_norm > max_norm
+        x_norm_outside = x_norm_outside.expand_as(x)
+
+        x = torch.where(x_norm_outside, x_unit*max_norm, x)
+
+        return x.view(x_shape)
+
+    @torch.enable_grad()
+    def attack(self, x, attack_targets, gt_labels, topk_coefs):
+        """
+        For a given set of topk coefficients, this function computes
+        best energy attack in the given number of iterations and configuration
+
+        x: [B, C, H, W] [0-1 for colors]
+        attack_targets: [B, K] (long)
+        gt_labels: [B] (long)
+        topk_coefs: [B] (floats)
+        """
+
+        topk_coefs = topk_coefs.clone()
+        K = attack_targets.shape[-1]
+
+        x_perturbation = nn.Parameter(torch.randn(x.shape, 
+                                                 device=x.device)*2e-3)
+
+        with torch.no_grad():
+            prediction_logits_0, prediction_feats_0 \
+                = self.model(x, return_features=True)
+
+        best_perturbations = torch.zeros_like(x) # [B, 3, H, W]
+        has_successful_attack = torch.zeros(x.shape[0], dtype=torch.long, device=x.device) # [B]
+        best_energy = torch.full((x.shape[0],), float('inf'), device=x.device) # [B]
+
+        pbar = tqdm(range(self.iterations))
+
+        for i in pbar:
+
+            if i == self.config.opt_warmup_its:
+                # Reset optimizer state
+                optimizer = self.optimizer([x_perturbation], lr=self.lr)
+
+            x_perturbed = x + x_perturbation#self.surject_perturbation(x_perturbation)
+            prediction_logits, prediction_feats = self.model(x_perturbed, return_features=True)
+            
+            pred_classes = prediction_logits.argsort(dim=-1, descending=True) # [B, C]
+            attack_successful = (pred_classes[:, :K] == attack_targets).all(dim=-1) # [B]
+            attack_energy = x_perturbation.flatten(1).norm(dim=-1) # [B]
+
+            attack_improved = attack_successful & (attack_energy <= best_energy)
+
+            best_perturbations[attack_improved] = x_perturbation[attack_improved]
+            has_successful_attack[attack_improved] = True
+            best_energy[attack_improved] = attack_energy[attack_improved]
+
+            loss = self.loss(logits_pred=prediction_logits,
+                             feats_pred=prediction_feats, 
+                             feats_pred_0=prediction_feats_0,
+                             attack_targets=attack_targets, 
+                             model=self.model, **precomputed_state)
+            
+            loss = loss * topk_coefs
+
+            loss = loss.sum()
+            
+            pbar.set_description(f"Loss: {loss.item():.3f}")
+            
+            loss = loss + x_perturbation.flatten(1).square().sum()
+            
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            # If we were successfull let's start taking the norm down
+            topk_coefs[attack_improved] *= 0.75
+
+            # Project perturbation to be within image limits
+            with torch.no_grad():
+                x_perturbed = x + x_perturbation
+                x_perturbed = x_perturbed.clamp_(min=0., max=1.)
+
+                x_perturbation.data = x_perturbed - x
+            
+        x_perturbed_best = x + best_perturbations
+        prediction_logits, prediction_feats = self.model(x_perturbed_best, return_features=True)
+
+        if self.config.dump_plots:
+            if os.path.isdir(self.config.plot_out):
+                shutil.rmtree(self.config.plot_out)
+
+            if has_successful_attack.any():
+                def dump_random_map():
+                    os.makedirs(self.config.plot_out, exist_ok=True)
+
+                    # selected_idx = best_energy.argmin()
+                    successful_idxs = has_successful_attack.nonzero()[:, 0]
+
+                    if self.config.plot_idx == "find":
+                        selected_idx = successful_idxs[torch.randperm(len(successful_idxs))[0]]
+                        # selected_idx = best_energy.argmin()
+                    else:
+                        selected_idx = int(self.config.plot_idx)
+
+                    print ("Selected idx", selected_idx)
+
+                    top_classes = prediction_logits_0[selected_idx].argsort(dim=-1, descending=True)
+                    attack_targets_selected = attack_targets[selected_idx]
+
+                    def imgnet_names(idxs):
+                        return [imgnet_idx_to_name[int(idx)].split(",")[0] for idx in idxs]
+                                    
+                    top_class_names = imgnet_names(top_classes)[:K]
+                    attack_targets_selected_names = imgnet_names(attack_targets_selected)
+
+                    def plot_attn_map(attn_map):
+                        attn_map = attn_map[0].mean(dim=0)[1:] # [196] get class tokens
+                        attn_map = attn_map.view(14, 14)
+                        attn_map = F.interpolate(
+                            attn_map[None, None],
+                            x.shape[-2:],
+                            mode="bilinear"
+                        ).view(x.shape[-2:])
+
+                        plt.imshow(attn_map.detach().cpu(), alpha=0.5)
+
+                    plt.figure()
+                    plt.imshow(x[selected_idx].permute(1,2,0).flip(dims=(-1,)).detach().cpu())
+                    plt.axis("off")
+                    plt.savefig(f"{self.config.plot_out}/clean_image.png", bbox_inches="tight", pad_inches=0)
+
+                    plt.figure()
+                    plt.imshow(x_perturbed_best[selected_idx].permute(1,2,0).flip(dims=(-1,)).detach().cpu())
+                    plt.axis("off")
+                    plt.savefig(f"{self.config.plot_out}/perturbed_image.png", bbox_inches="tight", pad_inches=0)
+
+                    plt.figure()
+                    plt.imshow(best_perturbations[selected_idx].mean(dim=0).abs().detach().cpu(), cmap="hot")
+                    plt.colorbar()
+                    plt.savefig(f"{self.config.plot_out}/perturbation.png", bbox_inches="tight")
+
+                    if isinstance(self.model, MMPretrainVisualTransformerWrapper):
+                        attn_maps_clean = self.model.get_attention_maps(x)[-1][selected_idx]
+                        attn_maps_attacked = self.model.get_attention_maps(x_perturbed_best)[-1][selected_idx]
+
+                        plt.figure()
+                        plt.imshow(x[selected_idx].permute(1,2,0).flip(dims=(-1,)).detach().cpu())
+                        plot_attn_map(attn_maps_clean)
+                        plt.axis("off")
+                        plt.savefig(f"{self.config.plot_out}/clean_map.png", bbox_inches="tight", pad_inches=0)
+
+                        plt.figure()
+                        plt.imshow(x[selected_idx].permute(1,2,0).flip(dims=(-1,)).detach().cpu())
+                        plot_attn_map(attn_maps_attacked)
+                        plt.axis("off")
+                        plt.savefig(f"{self.config.plot_out}/attacked_map.png", bbox_inches="tight", pad_inches=0)
+
+                    with open(f'{self.config.plot_out}/clean_classes_names.txt', 'w') as f:
+                        f.write(", ".join(top_class_names))
+
+                    with open(f'{self.config.plot_out}/attack_targets_names.txt', 'w') as f:
+                        f.write(", ".join(attack_targets_selected_names))
+
+                    with open(f'{self.config.plot_out}/clean_classes_names.txt', 'w') as f:
+                        f.write(", ".join(top_class_names))
+
+                    with open(f'{self.config.plot_out}/selected_idx.txt', 'w') as f:
+                        if isinstance(selected_idx, torch.Tensor):
+                            selected_idx = selected_idx.item()
+
+                        f.write(str(selected_idx))
+
+                    with open(f'{self.config.plot_out}/energy.txt', 'w') as f:
+                        f.write(str(best_energy[selected_idx].item()))
+                    
+                    C = prediction_logits_0.shape[-1]
+                    class_idxs = torch.arange(C) + 1
+                    clean_probs = prediction_logits_0[selected_idx].detach().cpu().softmax(dim=-1)
+                    attacked_probs = prediction_logits[selected_idx].detach().cpu().softmax(dim=-1)
+
+                    def label_classes(bars):
+                        adjusted_heights = {}
+                        for i, cls_idx in enumerate(attack_targets_selected.tolist()):
+                            bar = bars[cls_idx]
+                            height = bar.get_height()
+                            ann_x = bar.get_x() + bar.get_width()
+
+                            rotation = 90
+                            font_size = 10
+
+                            max_neighboring_height = -1
+                            for other_cls_idx in attack_targets_selected.tolist():
+                                if abs(cls_idx - other_cls_idx) <= 40 and cls_idx != other_cls_idx:
+                                    if other_cls_idx in adjusted_heights and adjusted_heights[other_cls_idx] > max_neighboring_height:
+                                        max_neighboring_height = adjusted_heights[other_cls_idx]
+                            
+                            if max_neighboring_height > 0:
+                                height = max_neighboring_height + 0.05
+
+                            adjusted_heights[cls_idx] = height
+
+                            plt.text(ann_x, height, f"[{i}]", rotation=rotation,
+                                    ha='center', va='bottom', fontsize=font_size, color='red')#.get_bbox_patch().get_height()
+                            
+
+                    plt.figure()
+                    bars_clean = plt.bar(class_idxs, clean_probs, width=4)
+                    plt.ylim(0,1)
+                    label_classes(bars_clean)
+                    plt.savefig(f"{self.config.plot_out}/clean_probs.png", bbox_inches="tight", pad_inches=0)
+
+                    plt.figure()
+                    bars_attacked = plt.bar(class_idxs, attacked_probs, width=4)
+                    plt.ylim(0,1)
+                    label_classes(bars_attacked)
+                    plt.savefig(f"{self.config.plot_out}/attacked_probs.png", bbox_inches="tight", pad_inches=0)
+
+                    print ("Idx", selected_idx)
+                    print (best_energy[selected_idx])
+                    print ("Finished plotting")
+
+                dump_random_map()
+                import sys
+                sys.exit(1)
+                print ("Dumped attention map")
+
+
+        return prediction_logits, best_perturbations, best_energy
+    
+    def forward(self, x, attack_targets, gt_labels):
+        """
+        This function is in charge of performing a binary search through
+        topk loss coefficients and running attacks on each.
+        """
+        B = x.shape[0]
+        device = x.device
+        topk_coefs_lower = torch.full((B,), fill_value=self.topk_loss_coef_lower,
+                                      device=device, dtype=torch.float)
+        
+        topk_coefs_upper = torch.full((B,), fill_value=self.topk_loss_coef_upper,
+                                      device=device, dtype=torch.float)
+        
+        best_perturbations = torch.zeros_like(x) # [B, 3, H, W]
+        best_energy = torch.full((B,), float('inf'), device=device) # [B]
+        best_prediction_logits = None
+
+        for search_step_i in range(self.binary_search_steps):
+            if x.device.index is None or x.device.index == 0:
+                print ("Running binary search step", search_step_i + 1)
+
+            current_topk_coefs = (topk_coefs_lower + topk_coefs_upper) / 2
+            current_logits, current_perturbations, current_energy = \
+                self.attack(x, attack_targets, gt_labels, current_topk_coefs)
+            
+            current_attack_suceeded = ~torch.isinf(current_energy)
+            
+            update_mask = current_energy < best_energy
+
+            best_perturbations[update_mask] = current_perturbations[update_mask]
+            best_energy[update_mask] = current_energy[update_mask]
+            
+            if best_prediction_logits is None:
+                best_prediction_logits = current_logits.clone()
+            else:
+                best_prediction_logits[update_mask] = current_logits[update_mask]
+
+            # If we fail to attack, we must increase our topk coef
+            topk_coefs_lower[~current_attack_suceeded] = current_topk_coefs[~current_attack_suceeded]
+
+            # If we succeed, we must lower to seek a more frugal attack
+            topk_coefs_upper[current_attack_suceeded] = current_topk_coefs[current_attack_suceeded]
+
+            idist.barrier()
+            
+        return best_prediction_logits, best_perturbations

modelguidedattacks/losses/__init__.py

@@ -0,0 +1,4 @@
+from .boilerplate import BoilerplateLoss
+from .cw_extension import CWExtensionLoss
+from .cvx_proj import CVXProjLoss
+from .adversarial_distillation.ad_loss import AdversarialDistillationLoss

modelguidedattacks/losses/_qp_solver_patch.py

@@ -0,0 +1,170 @@
+import qpth
+from qpth.solvers.pdipm import batch as pdipm_b
+from qpth.solvers.pdipm.batch import *
+
+def reduce_stats(z):
+    return z[~z.isnan()].median()
+
+def forward(Q, p, G, h, A, b, Q_LU, S_LU, R, eps=1e-12, verbose=0, notImprovedLim=3,
+            maxIter=20, solver=KKTSolvers.LU_PARTIAL):
+    """
+    Q_LU, S_LU, R = pre_factor_kkt(Q, G, A)
+    """
+    nineq, nz, neq, nBatch = get_sizes(G, A)
+
+    # Find initial values
+    if solver == KKTSolvers.LU_FULL:
+        D = torch.eye(nineq).repeat(nBatch, 1, 1).type_as(Q)
+        x, s, z, y = factor_solve_kkt(
+            Q, D, G, A, p,
+            torch.zeros(nBatch, nineq).type_as(Q),
+            -h, -b if b is not None else None)
+    elif solver == KKTSolvers.LU_PARTIAL:
+        d = torch.ones(nBatch, nineq).type_as(Q)
+        factor_kkt(S_LU, R, d)
+        x, s, z, y = solve_kkt(
+            Q_LU, d, G, A, S_LU,
+            p, torch.zeros(nBatch, nineq).type_as(Q),
+            -h, -b if neq > 0 else None)
+    elif solver == KKTSolvers.IR_UNOPT:
+        D = torch.eye(nineq).repeat(nBatch, 1, 1).type_as(Q)
+        x, s, z, y = solve_kkt_ir(
+            Q, D, G, A, p,
+            torch.zeros(nBatch, nineq).type_as(Q),
+            -h, -b if b is not None else None)
+    else:
+        assert False
+
+    # Make all of the slack variables >= 1.
+    M = torch.min(s, 1)[0]
+    M = M.view(M.size(0), 1).repeat(1, nineq)
+    I = M < 0
+    s[I] -= M[I] - 1
+
+    # Make all of the inequality dual variables >= 1.
+    M = torch.min(z, 1)[0]
+    M = M.view(M.size(0), 1).repeat(1, nineq)
+    I = M < 0
+    z[I] -= M[I] - 1
+
+    best = {'resids': None, 'x': None, 'z': None, 's': None, 'y': None}
+    nNotImproved = 0
+
+    for i in range(maxIter):
+        # affine scaling direction
+        rx = (torch.bmm(y.unsqueeze(1), A).squeeze(1) if neq > 0 else 0.) + \
+            torch.bmm(z.unsqueeze(1), G).squeeze(1) + \
+            torch.bmm(x.unsqueeze(1), Q.transpose(1, 2)).squeeze(1) + \
+            p
+        rs = z
+        rz = torch.bmm(x.unsqueeze(1), G.transpose(1, 2)).squeeze(1) + s - h
+        ry = torch.bmm(x.unsqueeze(1), A.transpose(
+            1, 2)).squeeze(1) - b if neq > 0 else 0.0
+        mu = torch.abs((s * z).sum(1).squeeze() / nineq)
+        z_resid = torch.norm(rz, 2, 1).squeeze()
+        y_resid = torch.norm(ry, 2, 1).squeeze() if neq > 0 else 0
+        pri_resid = y_resid + z_resid
+        dual_resid = torch.norm(rx, 2, 1).squeeze()
+        resids = pri_resid + dual_resid + nineq * mu
+
+        d = z / s
+        try:
+            factor_kkt(S_LU, R, d)
+        except:
+            return best['x'], best['y'], best['z'], best['s']
+
+        if verbose == 1:
+            print('iter: {}, pri_resid: {:.5e}, dual_resid: {:.5e}, mu: {:.5e}'.format(
+                i, reduce_stats(pri_resid), reduce_stats(dual_resid), reduce_stats(mu)))
+        if best['resids'] is None:
+            best['resids'] = resids
+            best['x'] = x.clone()
+            best['z'] = z.clone()
+            best['s'] = s.clone()
+            best['y'] = y.clone() if y is not None else None
+            nNotImproved = 0
+        else:
+            I = resids < best['resids']
+            if I.sum() > 0:
+                nNotImproved = 0
+            else:
+                nNotImproved += 1
+            I_nz = I.repeat(nz, 1).t()
+            I_nineq = I.repeat(nineq, 1).t()
+            best['resids'][I] = resids[I]
+            best['x'][I_nz] = x[I_nz]
+            best['z'][I_nineq] = z[I_nineq]
+            best['s'][I_nineq] = s[I_nineq]
+            if neq > 0:
+                I_neq = I.repeat(neq, 1).t()
+                best['y'][I_neq] = y[I_neq]
+        if nNotImproved == notImprovedLim or reduce_stats(pri_resid) < eps or mu.min() > 1e32:
+            if best['resids'].max() > 1. and verbose >= 0:
+                print(INACC_ERR)
+            return best['x'], best['y'], best['z'], best['s']
+
+        if solver == KKTSolvers.LU_FULL:
+            D = bdiag(d)
+            dx_aff, ds_aff, dz_aff, dy_aff = factor_solve_kkt(
+                Q, D, G, A, rx, rs, rz, ry)
+        elif solver == KKTSolvers.LU_PARTIAL:
+            dx_aff, ds_aff, dz_aff, dy_aff = solve_kkt(
+                Q_LU, d, G, A, S_LU, rx, rs, rz, ry)
+        elif solver == KKTSolvers.IR_UNOPT:
+            D = bdiag(d)
+            dx_aff, ds_aff, dz_aff, dy_aff = solve_kkt_ir(
+                Q, D, G, A, rx, rs, rz, ry)
+        else:
+            assert False
+
+        # compute centering directions
+        alpha = torch.min(torch.min(get_step(z, dz_aff),
+                                    get_step(s, ds_aff)),
+                          torch.ones(nBatch).type_as(Q))
+        alpha_nineq = alpha.repeat(nineq, 1).t()
+        t1 = s + alpha_nineq * ds_aff
+        t2 = z + alpha_nineq * dz_aff
+        t3 = torch.sum(t1 * t2, 1).squeeze()
+        t4 = torch.sum(s * z, 1).squeeze()
+        sig = (t3 / t4)**3
+
+        rx = torch.zeros(nBatch, nz).type_as(Q)
+        rs = ((-mu * sig).repeat(nineq, 1).t() + ds_aff * dz_aff) / s
+        rz = torch.zeros(nBatch, nineq).type_as(Q)
+        ry = torch.zeros(nBatch, neq).type_as(Q) if neq > 0 else torch.Tensor()
+
+        if solver == KKTSolvers.LU_FULL:
+            D = bdiag(d)
+            dx_cor, ds_cor, dz_cor, dy_cor = factor_solve_kkt(
+                Q, D, G, A, rx, rs, rz, ry)
+        elif solver == KKTSolvers.LU_PARTIAL:
+            dx_cor, ds_cor, dz_cor, dy_cor = solve_kkt(
+                Q_LU, d, G, A, S_LU, rx, rs, rz, ry)
+        elif solver == KKTSolvers.IR_UNOPT:
+            D = bdiag(d)
+            dx_cor, ds_cor, dz_cor, dy_cor = solve_kkt_ir(
+                Q, D, G, A, rx, rs, rz, ry)
+        else:
+            assert False
+
+        dx = dx_aff + dx_cor
+        ds = ds_aff + ds_cor
+        dz = dz_aff + dz_cor
+        dy = dy_aff + dy_cor if neq > 0 else None
+        alpha = torch.min(0.999 * torch.min(get_step(z, dz),
+                                            get_step(s, ds)),
+                          torch.ones(nBatch).type_as(Q))
+        alpha_nineq = alpha.repeat(nineq, 1).t()
+        alpha_neq = alpha.repeat(neq, 1).t() if neq > 0 else None
+        alpha_nz = alpha.repeat(nz, 1).t()
+
+        x += alpha_nz * dx
+        s += alpha_nineq * ds
+        z += alpha_nineq * dz
+        y = y + alpha_neq * dy if neq > 0 else None
+
+    if best['resids'].max() > 1. and verbose >= 0:
+        print(INACC_ERR)
+    return best['x'], best['y'], best['z'], best['s']
+
+pdipm_b.forward = forward

modelguidedattacks/losses/adversarial_distillation/ad_loss.py

@@ -0,0 +1,38 @@
+import torch
+from torch import nn
+from .adversarial_distribution import AD_Distribution
+
+class AdversarialDistillationLoss(nn.Module):
+    def __init__(self, confidence=0, alpha=10, beta=0.3):
+        super().__init__()
+
+        self.alpha = alpha
+        self.beta = beta
+
+        self.distri_generator = AD_Distribution(simi_name='glove', 
+                                                alpha=self.alpha, beta=self.beta)
+        
+        self.kl = nn.KLDivLoss(reduction='none')
+        self.logsoftmax = nn.LogSoftmax(dim=-1)
+
+    def precompute(self, attack_targets, gt_labels, config): 
+        device = attack_targets.device 
+
+        target_distribution = self.distri_generator.generate_distribution(gt_labels.cpu(), attack_targets.cpu())
+        target_distribution = torch.from_numpy(target_distribution).float().to(device)
+
+        K = attack_targets.shape[-1]
+        target_distribution_topk = target_distribution.argsort(dim=-1, descending=True)[:, :K]
+
+        assert (target_distribution_topk == attack_targets).all()
+
+        return {
+            "ad_distribution": target_distribution
+        }
+    
+    def forward(self, logits_pred, feats_pred, feats_pred_0, attack_targets, model, ad_distribution, **kwargs):
+        log_logits = self.logsoftmax(logits_pred)
+        loss_kl = self.kl(log_logits, ad_distribution)
+        loss_kl = torch.sum(loss_kl, dim = -1)
+
+        return loss_kl

modelguidedattacks/losses/adversarial_distillation/adversarial_distribution.py

@@ -0,0 +1,72 @@
+import os
+import numpy as np
+from .glove_simi import generate_glove, AD_DIRECTORY
+
+class AD_Distribution():
+    def __init__(self, simi_name, alpha, beta):
+        print('using ........ ', simi_name, ' .... knowledge')
+        path_simi_name = os.path.join(AD_DIRECTORY, 'imagenet_cos_similarity')
+        file_simi_name = path_simi_name +'_'+ simi_name + '.npy'
+        if os.path.exists(file_simi_name):
+            self.cos_similarity = np.load(file_simi_name)
+            print(simi_name+" cos_similarity loaded")
+        else:
+            self.cos_similarity = self.generate_similarity(simi_name)
+        
+        self.alpha = alpha
+        self.beta  = beta
+    
+    def generate_similarity(self,simi_name):
+        if simi_name == 'glove':
+            similarity = generate_glove()
+        else:
+            print(simi_name + 'not implemented yet')
+        return similarity
+
+    def generate_distribution(self, gt_label, target):
+        distribution=[]
+
+        for i in range(len(target)):
+            distri = self.single_distribution_build(i, target[i], gt_label[i])
+            distribution.append(distri)
+
+        distribution = np.array(distribution)
+        return distribution
+    
+    def single_distribution_build(self,index, target_id, gt_id):
+        if target_id.shape == ():
+            target_id = np.array([target_id])
+
+        simil_logits = np.zeros(self.cos_similarity[target_id[0],:].shape)
+        for i in range(len(target_id)):
+            simil_logits += self.cos_similarity[target_id[i],:]
+            
+        simil_logits = (simil_logits)/ len(target_id)
+        logit_value = self.alpha
+
+        for i in range(len(target_id)):
+            simil_logits[target_id[i]] = logit_value
+            logit_value = logit_value - self.beta 
+            
+        if not self.check_oreder_target_no_groundtruth(simil_logits, target_id):
+            print('fail to generate distribution for index: ', index) 
+            
+        logits = self.softmax(simil_logits)
+        return logits
+
+    def check_oreder_target_no_groundtruth(self, probs, target_id):
+        sort_labels = np.argsort(probs)
+        cnt = 0
+        for i in range(len(target_id)):
+            if target_id[-(i+1)] == sort_labels[-(len(target_id)-i)]: 
+                cnt +=1
+        if (cnt == len(target_id)):
+            return True
+        else:
+            return False
+
+    def softmax(self,logits):
+        
+        prob=np.exp(logits) / np.sum(np.exp(logits))
+        return prob
+

modelguidedattacks/losses/adversarial_distillation/glove.py

@@ -0,0 +1,58 @@
+import torch
+from tqdm import tqdm
+
+class GloVe():
+
+   def __init__(self, file_path):
+       self.dimension = None
+       self.embedding = dict()
+       with open(file_path, 'r') as f:
+           for line in tqdm(f.readlines()):
+               strs = line.strip().split(' ')
+               word = strs[0]
+               vector = torch.FloatTensor(list(map(float, strs[1:])))
+               self.embedding[word] = vector
+               if self.dimension is None:
+                   self.dimension = len(vector)
+
+   def _fix_word(self, word):
+       terms = word.replace('_', ' ').split(' ')
+       ret = self.zeros()
+       cnt = 0
+       for term in terms:
+           v = self.embedding.get(term)
+           if v is None:
+               subterms = term.split('-')
+               subterm_sum = self.zeros()
+               subterm_cnt = 0
+               for subterm in subterms:
+                   subv = self.embedding.get(subterm)
+                   if subv is not None:
+                       subterm_sum += subv
+                       subterm_cnt += 1
+               if subterm_cnt > 0:
+                   v = subterm_sum / subterm_cnt
+           if v is not None:
+               ret += v
+               cnt += 1
+       return ret / cnt if cnt > 0 else None
+
+   def __getitem__(self, words):
+       if type(words) is str:
+           words = [words]
+       ret = self.zeros()
+       cnt = 0
+       for word in words:
+           v = self.embedding.get(word)
+           if v is None:
+               v = self._fix_word(word)
+           if v is not None:
+               ret += v
+               cnt += 1
+       if cnt > 0:
+           return ret / cnt
+       else:
+           return self.zeros()
+
+   def zeros(self):
+       return torch.zeros(self.dimension)

modelguidedattacks/losses/adversarial_distillation/glove_simi.py

@@ -0,0 +1,126 @@
+import os
+import numpy as np
+import requests
+import zipfile
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import optim
+from torch.autograd import Variable
+from torchvision import datasets, transforms
+import torchvision
+from PIL import Image
+from .glove import GloVe
+from tqdm import tqdm
+
+AD_DIRECTORY = os.path.dirname(__file__)
+
+def obtain_vector(inputs, glove):
+    vector_im = glove.embedding.get(inputs)
+    if vector_im is None:
+          vector_im = glove.embedding.get(inputs.lower())
+    if vector_im is None:
+          vector_im = glove.embedding.get(inputs.title())
+    if vector_im is None:
+          vector_im = glove.embedding.get(inputs.upper())
+    return vector_im
+
+def generate_glove():
+    print("Generating glove similarity...")
+    #download glove file
+    os.makedirs("./knowledge", exist_ok=True)
+    glove_file = './knowledge/glove.840B.300d.txt'
+    if not os.path.exists(glove_file):
+        print("Downloading glove files...")
+        print("")
+        print("Gonna take a while")
+        print("")
+        url_path = "http://nlp.stanford.edu/data/glove.840B.300d.zip"
+        r = requests.get(url_path)
+        with open("./knowledge/glove.840B.300d.zip","wb") as f:
+            f.write(r.content)
+        filename = './knowledge/glove.840B.300d.zip'
+        fz = zipfile.ZipFile(filename, 'r')
+        for file in fz.namelist():
+            fz.extract(file, './knowledge/.')
+        if os.path.exists(filename):
+              os.remove(filename)
+
+    glove = GloVe('./knowledge/glove.840B.300d.txt')
+    filepath = os.path.join(AD_DIRECTORY, "label_name.txt")
+    vec_list = []
+    vec_list_np = []
+    cos_similarity = np.zeros((1000,1000))
+
+    index = 0
+    
+    #the labels could be a word or a phrase with multi words
+    #we also tested on average of every words
+    #But we assume the last word should be more important, so in our final version
+    #we assign a higher weight to last word in a phrase and average the fornt words
+
+    #w2v for last word of multi-words
+    for line in tqdm(open(filepath)):
+      a = line.strip('\n')
+
+      b = a.split(',')
+      cnt = 0
+      vector = torch.zeros(300)
+      vec_front = torch.zeros(300)
+      vec_b_average = torch.zeros(300)
+      cnt_b = 0
+      for i in range(len(b)):
+            b[i] = b[i].lstrip()
+            c = b[i].split(' ')
+            if obtain_vector(c[-1], glove) is not None:
+              vec_b_average += obtain_vector(c[-1], glove)
+              cnt_b += 1
+      if cnt_b == 0:
+            print('index ', index,' generatint word_vector failure')
+            continue
+      vec_b_average = vec_b_average / cnt_b
+
+      for i in range(len(b)):
+            b[i] = b[i].lstrip()
+            c = b[i].split(' ')
+            cnt_f = 0
+            for j in range(len(c) - 1):
+                  if obtain_vector(c[j], glove) is not None:
+                        vec_front += obtain_vector(c[j], glove)
+                        cnt_f += 1
+            if obtain_vector(c[-1], glove) is not None:
+                  vec_back =obtain_vector(c[-1], glove)
+            else:
+                  vec_back = vec_b_average
+            if cnt_f == 0:
+                  vector += vec_back
+            else:
+                  vector += (vec_front / cnt_f )* 0.1 + vec_back * 0.9
+            cnt += 1
+
+      vector = torch.div(vector,cnt) 
+  
+      vec_list_np.append(np.array(vector))
+      vec_list.append(vector)
+      index += 1
+
+
+    vec_list_np_stacked = np.stack(vec_list_np)
+    vec_list_torch = torch.from_numpy(vec_list_np_stacked)
+
+    cos_similarity = F.cosine_similarity(vec_list_torch[None, :], vec_list_torch[:, None], dim=-1)
+    cos_similarity = cos_similarity.numpy()
+
+#     np.save('./knowledge/golve_vec_list', np.array(vec_list_np))
+#     for i in range(len(vec_list)):
+#       for j in range(len(vec_list)):
+#         cos_similarity[i,j] = F.cosine_similarity(vec_list[i], vec_list[j],dim=0).type(torch.half)
+#         if i != j:
+#           cos_similarity[i,j] = cos_similarity[i,j]
+#     cos_similarity = np.array(cos_similarity)
+    
+    
+    np.save(os.path.join(AD_DIRECTORY, "imagenet_cos_similarity_glove"), cos_similarity)
+    print("Glove cos_similarity finished")
+    return cos_similarity

modelguidedattacks/losses/boilerplate.py

@@ -0,0 +1,59 @@
+import torch
+from torch import nn
+
+def generalized_mean(x, p, dim):
+    x_type = x.dtype
+    x = x.to(torch.double)
+    x = x**p
+    x = x.mean(dim=dim)
+    x = x**(1/p)
+    return x.to(x_type)
+
+def surject_to_positive(x, c=5):
+    assert x.min() >= -1
+    assert x.max() <= 1
+
+    return c + c * x
+
+def surject_from_positive(x, c=5):
+    return (x - c) / c
+
+class BoilerplateLoss(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.p = 9
+
+    def forward(self, y_pred, y_attack, **kwargs):
+        y_pred = y_pred.softmax(dim=-1)
+
+        C = y_pred.shape[1]
+        K = y_attack.shape[1]
+        desired_mask = torch.zeros_like(y_pred, dtype=torch.bool)
+        desired_mask.scatter_(dim=1, index=y_attack, 
+                              src=torch.ones_like(y_attack, dtype=torch.bool))
+        
+        y_not_in_attack = (~desired_mask).nonzero()[:, 1].view(-1, C - K)
+
+        y_pred_in_attack = torch.gather(y_pred, dim=1, index=y_attack)
+        y_pred_not_in_attack = torch.gather(y_pred, dim=1, index=y_not_in_attack)
+
+        y_pred_in_attack_min = y_pred_in_attack.min(dim=-1).values #generalized_mean(y_pred_in_attack, -self.p, dim=1)
+        y_pred_not_in_attack_max = y_pred_not_in_attack.max(dim=-1).values #generalized_mean(y_pred_not_in_attack, self.p, dim=1)
+
+        macro_loss = (y_pred_not_in_attack_max - y_pred_in_attack_min)
+        sorting_loss = y_pred_in_attack.diff(dim=-1)
+
+        # Surject sorting_loss to positive domain, since it goes [-1,1] we can just shift by 1
+        sorting_loss = surject_to_positive(sorting_loss)
+        sorting_loss = generalized_mean(sorting_loss, p=9, dim=-1)
+
+        # Surject back
+        sorting_loss = surject_from_positive(sorting_loss)
+
+        catted_loss = torch.stack([macro_loss, sorting_loss], dim=-1)
+        catted_loss_pos = surject_to_positive(catted_loss)
+        
+        final_loss_pos = generalized_mean(catted_loss_pos, p=10, dim=-1)
+        final_loss = surject_from_positive(final_loss_pos)
+
+        return final_loss

modelguidedattacks/losses/cvx_proj.py

@@ -0,0 +1,108 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+import cvxpy as cp
+import qpth
+from . import _qp_solver_patch
+
+def solve_qp(Q, P, G, H):
+    B = Q.shape[0]
+
+    if B == 1:
+        # Batch size of 1 has weird instabilities
+        # I imagine there is a .squeeze() or something inside the QP solver's code
+        # that messes up broadcasting dimensions when batch dimension is 1 so let's
+        # artificially make 2 solutions when we need 1
+
+        Q = Q.expand(2, -1, -1)
+        P = P.expand(2, -1)
+        G = G.expand(2, -1, -1)
+        H = H.expand(2, -1)
+
+    e = torch.empty(0, device=Q.device)
+    z_sol = qpth.qp.QPFunction(verbose=-1, eps=1e-2, check_Q_spd=False)(Q, P, G, H, e, e)
+
+    if B == 1:
+        z_sol = z_sol[:1]
+
+    return z_sol
+
+class CVXProjLoss(nn.Module):
+    def __init__(self, confidence=0):
+        super().__init__()
+        self.confidence = confidence
+
+    def precompute(self, attack_targets, gt_labels, config):
+        return {
+            "margin": config.cvx_proj_margin
+        }
+
+    def forward(self, logits_pred, feats_pred, feats_pred_0, attack_targets, model, margin, **kwargs):
+        device = logits_pred.device
+        head_W, head_bias = model.head_matrices()
+
+        num_feats = head_W.shape[1]
+        num_classes = head_W.shape[0]
+
+        K = attack_targets.shape[-1]
+        B = logits_pred.shape[0]
+        
+        # Start with all classes should be less than smallest attack target
+        D = -torch.eye(num_classes, device=device)[None].repeat(B, 1, 1) # [B, C, C]
+        attack_targets_write = attack_targets[:, -1][:, None, None].expand(-1, D.shape[1], -1)
+        D.scatter_(dim=2, index=attack_targets_write, src=torch.ones(attack_targets_write.shape, device=device))
+
+        # Clear out the constraint row for each item in the attack targets
+        attack_targets_clear = attack_targets[:, :, None].expand(-1, -1, D.shape[-1])
+        D.scatter_(dim=1, index=attack_targets_clear, src=torch.zeros(attack_targets_clear.shape, device=device))
+
+        batch_inds = torch.arange(B, device=device)[:, None].expand(-1, K - 1)
+        attack_targets_pos = attack_targets[:, :-1] # [B, K-1]
+        attack_targets_neg = attack_targets[:, 1:] # [B, K-1]
+
+        attack_targets_neg_inds = torch.stack((
+            batch_inds,
+            attack_targets_neg,
+            attack_targets_neg
+        ), dim=0) # [3, B, K - 1]
+        attack_targets_neg_inds = attack_targets_neg_inds.view(3, -1)
+
+        D[attack_targets_neg_inds[0], attack_targets_neg_inds[1], attack_targets_neg_inds[2]] = -1
+
+        attack_targets_pos_inds = torch.stack((
+            batch_inds,
+            attack_targets_neg,
+            attack_targets_pos
+        ), dim=0) # [3, B, K - 1]
+
+        D[attack_targets_pos_inds[0], attack_targets_pos_inds[1], attack_targets_pos_inds[2]] = 1
+
+        A = head_W
+        b = head_bias
+
+        Q = 2*torch.eye(feats_pred.shape[1], device=device)[None].expand(B, -1, -1)
+
+        # We want the solution features to be as close as possible
+        # to the current features but also head on the direction of 
+        # the smallest possible perturbation from the initial predicted
+        # features
+        anchor_feats = feats_pred
+
+        P = -2*anchor_feats.expand(B, -1)
+
+        G = -D@A
+        H = -(margin - D @ b)
+
+        # Constraints are indexed by smaller logit
+        # First attack target isn't smaller than any logit, so its 
+        # constraint index is redundant, but we keep it for easier parallelization
+        # Make this constraint all 0s
+        zero_inds = attack_targets[:, 0:1] # [B, 1]
+        H.scatter_(dim=1, index=zero_inds, src=torch.zeros(zero_inds.shape, device=device))
+
+        z_sol = solve_qp(Q, P, G, H)
+
+        loss = (feats_pred - z_sol).square().sum(dim=-1)
+
+        # loss_check = self.forward_check(logits_pred, feats_pred, attack_targets, model, **kwargs)
+        return loss

modelguidedattacks/losses/cw_extension.py

@@ -0,0 +1,46 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+CARLINI_COEFF_UPPER = 1e10
+
+class CWExtensionLoss(nn.Module):
+    def __init__(self, confidence=0):
+        super().__init__()
+        self.confidence = confidence
+
+    def precompute(self, *args, **kwargs):
+        return {}
+
+    def forward(self, logits_pred, attack_targets,  **kwargs):
+        #orign cw attack loss
+        if attack_targets.dim() == 1:
+            mask_logits = F.one_hot(attack_targets, logits_pred.shape[1]).float()
+
+            real = (mask_logits * logits_pred).sum(dim=1)
+            other = ((1.0 - mask_logits) * logits_pred - (mask_logits * 10000.0)
+                    ).max(1)[0]
+            loss_cw = torch.clamp(other - real + self.confidence, min=0.)
+            return loss_cw
+            
+        #extended cw loss for topk attack tasks
+        else:
+            mask_logits = torch.zeros([logits_pred.shape[0], logits_pred.shape[1]], device=logits_pred.device)
+            min_values = torch.ones(attack_targets.shape[0], dtype=torch.float, device=logits_pred.device) * 1e10
+            loss_cw_topk = 0
+
+            for i in range(attack_targets.shape[1]):
+                other = ((1.0 - mask_logits) * logits_pred - (mask_logits * 10000.0)
+                    ).max(1)[0]
+
+
+                loss_cw_topk += torch.clamp(other - min_values + self.confidence, min=0.)
+                mask_logits[torch.arange(len(attack_targets)), attack_targets[:,i]] = 1
+                min_values = torch.min(logits_pred[torch.arange(len(attack_targets)), attack_targets[:,i]], min_values)
+
+            real =  min_values
+            other = ((1.0 - mask_logits) * logits_pred - (mask_logits * 10000.0)
+                    ).max(1)[0]
+            loss_cw_topk += torch.clamp(other - real + self.confidence, min=0.)
+            constant = attack_targets.shape[1]
+
+            return (loss_cw_topk / constant)

modelguidedattacks/losses/energy.py

@@ -0,0 +1,80 @@
+import torch
+from torch import nn
+from ignite.metrics import Loss
+from ignite.metrics.metric import reinit__is_reduced, sync_all_reduce
+from typing import Callable, cast, Dict, Sequence, Tuple, Union
+
+def get_correct_mask(y_pred, y_attack):
+    k = y_attack.shape[-1]
+
+    y_pred_indices = y_pred.argsort(dim=-1, descending=True) # [N, C]
+
+    correct = (y_pred_indices[:, :k] == y_attack).all(dim=-1)
+    return correct
+
+class EnergyLoss(Loss):
+    def __init__(self, loss_fn, reduction="mean", device = ...):
+        super().__init__(loss_fn, device=device)
+        self.reduction = reduction
+
+    @reinit__is_reduced
+    def reset(self) -> None:
+        self._sum = torch.tensor(0.0, device=self._device)
+        self._min = torch.tensor(torch.inf, device=self._device)
+        self._max = torch.tensor(0.0, device=self._device)
+        self._num_examples = 0
+
+    @reinit__is_reduced
+    def update(self, output: Sequence[Union[torch.Tensor, Dict]]) -> None:
+        if len(output) == 2:
+            y_pred, y = cast(Tuple[torch.Tensor, torch.Tensor], output)
+            kwargs: Dict = {}
+        else:
+            y_pred, y, kwargs = cast(Tuple[torch.Tensor, torch.Tensor, Dict], output)
+
+        sample_energies = self._loss_fn(y_pred, y, **kwargs).detach()
+
+        n = len(sample_energies)
+
+        if n > 0:
+            self._sum += sample_energies.sum()
+            self._min = torch.minimum(self._min, sample_energies.min())
+            self._max = torch.maximum(self._max, sample_energies.max())
+            self._num_examples += n
+
+    @sync_all_reduce("_sum", "_num_examples", "_min:MIN", "_max:MAX")
+    def compute(self) -> float:
+
+        if self.reduction == "mean":
+            if self._num_examples == 0:
+                return torch.inf
+            
+            return self._sum.item() / self._num_examples
+        elif self.reduction == "max":
+            if self._num_examples == 0:
+                return torch.nan
+    
+            return self._max.item()
+        elif self.reduction == "min":
+            if self._num_examples == 0:
+                return torch.inf
+            
+            return self._min.item()
+        else:
+            assert False
+
+class Energy(nn.Module):
+    def __init__(self, p="2") -> None:
+        super().__init__()
+        self.p = p
+
+    def forward(self, y_pred, y_attack, perturbations, **kwargs):
+        correct = get_correct_mask(y_pred, y_attack)
+
+        # Don't want to take into account perturbations of
+        # unsuccessful attacks
+
+        perturbations = perturbations[correct]
+        perturbations = perturbations.flatten(1)
+
+        return torch.linalg.vector_norm(perturbations, dim=-1, ord=self.p)

modelguidedattacks/metrics/topk_accuracy.py

@@ -0,0 +1,15 @@
+import torch
+from ignite.metrics import Accuracy, Loss
+from typing import Sequence
+
+class TopKAccuracy(Accuracy):
+    def update(self, output: Sequence[torch.Tensor], **kwargs) -> None:
+        y_pred, y_attack = output[0].detach(), output[1].detach()
+        k = y_attack.shape[-1]
+
+        y_pred_indices = y_pred.argsort(dim=-1, descending=True) # [N, C]
+
+        correct = (y_pred_indices[:, :k] == y_attack).all(dim=-1)
+
+        self._num_correct += torch.sum(correct).to(self._device)
+        self._num_examples += correct.shape[0]

modelguidedattacks/models.py

@@ -0,0 +1,32 @@
+from torchvision import models
+
+from modelguidedattacks.guides.instance_guide import InstanceGuide
+from modelguidedattacks.guides.unguided import Unguided
+from modelguidedattacks import losses
+
+from .cls_models.registry import get_model
+
+guide_model_registry = {
+    "instance_guided": InstanceGuide,
+    "unguided": Unguided
+}
+
+loss_registry = {
+    "cvxproj": losses.CVXProjLoss,
+    "cwk": losses.CWExtensionLoss,
+    "ad": losses.AdversarialDistillationLoss
+}
+
+def setup_model(config, device):
+    model = get_model(config.dataset, config.model, device)
+
+    kwargs = {}
+
+    if config.guide_model == "unguided":
+        kwargs["iterations"] = config.unguided_iterations
+        kwargs["lr"] = config.unguided_lr
+        kwargs["loss_fn"] = loss_registry[config.loss]
+        kwargs["binary_search_steps"] = config.binary_search_steps
+        kwargs["topk_loss_coef_upper"] = config.topk_loss_coef_upper
+
+    return guide_model_registry[config.guide_model](model, config, **kwargs)

modelguidedattacks/results.py

@@ -0,0 +1,261 @@
+import os
+import torch
+from pathlib import Path
+import numpy as np
+import copy
+from collections import OrderedDict
+
+config_parameter_keys = ["loss", "unguided_lr", "model", "k", "binary_search_steps",
+                        "unguided_iterations", "topk_loss_coef_upper", "seed",
+                        "opt_warmup_its", "cvx_proj_margin", 
+                        "topk_loss_coef_upper", "binary_search_steps"]
+
+def config_to_dict(config):
+    result_keys = config_parameter_keys
+        
+    result_dict = {
+        key: getattr(config, key) for key in result_keys
+    }
+
+    if hasattr(config, "cvx_proj_margin"):
+        result_dict["cvx_proj_margin"] = config.cvx_proj_margin
+    else:
+        result_dict["cvx_proj_margin"] = 0.2
+
+    return result_dict
+
+def load_all_results(load_min_max=False):
+    if not os.path.isdir("results_rebuttal"):
+        return []
+    
+    all_result_files = Path('results_rebuttal').rglob('*.save')
+
+    results_list = []
+    for result_file in all_result_files:
+        result = torch.load(result_file)
+        config = result["config"]
+        
+        result_dict = config_to_dict(config)
+
+        result_dict["ASR"] = result["ASR"]
+        result_dict["L1"] = result["L1 Energy"]
+        result_dict["L2"] = result["L2 Energy"]
+        result_dict["L_inf"] = result["L_inf Energy"]
+
+        if "L2 Energy Max" in result and load_min_max:
+            result_dict["L1 Max"] = result["L1 Energy Max"]
+            result_dict["L2 Max"] = result["L2 Energy Max"]
+            result_dict["L_inf Max"] = result["L_inf Energy Max"]
+
+            result_dict["L1 Min"] = result["L1 Energy Min"]
+            result_dict["L2 Min"] = result["L2 Energy Min"]
+            result_dict["L_inf Min"] = result["L_inf Energy Min"]
+
+        results_list.append(result_dict)
+
+    return results_list
+
+def close(target, eps=1e-5):
+    return lambda x: np.allclose(x, target, atol=eps)
+
+def eq(target):
+    if isinstance(target, float):
+        return close(target)
+    else:
+        return lambda x: x == target
+
+def gte(target):
+    return lambda x: float(x) >= target
+
+def lte(target):
+    return lambda x: float(x) <= target
+
+def in_set(target):
+    return lambda x: x in target
+
+def filter_from_config(config):
+    config_dict = config_to_dict(config)
+
+    filter = {
+        key: eq(val) for (key, val) in config_dict.items()
+    }
+
+    return filter
+
+def filter_results(filter, results_list, only_with_minmax=False):
+    filtered_results = []
+    for result in results_list:
+        pass_filter = True
+        for key, val in result.items():
+            if key not in filter:
+                continue
+
+            if not filter[key](val):
+                pass_filter = False
+                break
+
+        if only_with_minmax and "L2 Max" not in result:
+            continue
+
+        if pass_filter:
+            filtered_results.append(result)
+
+    return filtered_results
+
+def resolve_nonunique_filter(filter, results_list, include_failed=False):
+    filtered_results = filter_results(filter, results_list)
+
+    unique_parameters = []
+    # Find unique parameter sets for results
+    for result in filtered_results:
+        result_parameters = {param_key:result[param_key] for param_key in config_parameter_keys}
+
+        # Round to avoid floating pt imprecision from messing with set uniqueness checks
+        for key in result_parameters.keys():
+            if isinstance(result_parameters[key], float):
+                result_parameters[key] = round(result_parameters[key], 5)
+
+        del result_parameters["seed"]
+        unique_parameters.append(result_parameters)
+
+    # Only keep unique dicts
+    unique_parameters = [dict(y) for y in set(tuple(x.items()) for x in unique_parameters)]
+
+    best_metric = -np.Infinity
+    best_param_set = None
+    best_result_list = None
+    for param_set in unique_parameters:
+        # Perform another search
+        unique_filter = {
+            param_name: eq(param_value) for param_name, param_value in param_set.items()
+        }
+
+        filtered_results = filter_results(unique_filter, results_list)
+
+        assert len(filtered_results) == 5
+
+        asrs = [result["ASR"] for result in filtered_results]
+        l2_energies = [result["L2"] for result in filtered_results]
+
+        mean_asr = np.mean(np.array(asrs)[np.isfinite(asrs)])
+        mean_l2 = np.mean(np.array(l2_energies)[np.isfinite(l2_energies)])
+
+        # Arbitrary point in tradeoff curve
+        result_goodness = -mean_l2 + mean_asr * 100
+
+        if (mean_asr > 0 and mean_asr < 0.025) and not include_failed:
+            # Irrelevant result and associated energies
+            continue
+
+        if result_goodness > best_metric or (include_failed and best_param_set is None):
+            best_param_set = param_set
+            best_result_list = filtered_results
+            best_metric = result_goodness
+
+    return best_param_set, best_result_list
+
+def get_combined_results(filtered_results):
+    combined_results = {}
+
+    for result in filtered_results:
+        for key in result:
+            if key not in combined_results:
+                combined_results[key] = []
+
+            combined_results[key].append(result[key])
+
+    unique_runs = len(np.unique(combined_results["seed"]))
+    # assert len(combined_results["seed"]) == unique_runs
+
+    for key, val in list(combined_results.items()):
+        if key in ["ASR", "L1", "L2", "L_inf"]:
+            val = np.array(val)
+            combined_results[f"{key}_mean"] = np.mean(val[np.isfinite(val)])
+            combined_results[f"{key}_median"] = np.median(val[np.isfinite(val)])
+
+    # Coupled results
+    best_asr_idx = np.argmax(combined_results["ASR"])
+    best_asr = combined_results["ASR"][best_asr_idx]
+    best_l1 = combined_results["L1"][best_asr_idx]
+    best_l2 = combined_results["L2"][best_asr_idx]
+    best_linf = combined_results["L_inf"][best_asr_idx]
+
+    combined_results["ASR_best"] = best_asr
+    combined_results["L1_best"] = best_l1
+    combined_results["L2_best"] = best_l2
+    combined_results["L_inf_best"] = best_linf
+
+    worst_asr_idx = np.argmin(combined_results["ASR"])
+    worst_asr = combined_results["ASR"][worst_asr_idx]
+    worst_l1 = combined_results["L1"][worst_asr_idx]
+    worst_l2 = combined_results["L2"][worst_asr_idx]
+    worst_linf = combined_results["L_inf"][worst_asr_idx]
+
+    combined_results["ASR_worst"] = worst_asr
+    combined_results["L1_worst"] = worst_l1
+    combined_results["L2_worst"] = worst_l2
+    combined_results["L_inf_worst"] = worst_linf
+
+    return combined_results
+
+def build_full_results_dict(model_name="resnet50", verbose=False,
+                            all_k=[20, 15, 10, 5, 1], 
+                            all_num_iter=[60, 30],
+                            all_search_steps=[1, 9],
+                            all_methods=["cwk", "ad", "cvxproj"]):
+
+    if verbose:
+        print ("-" * 100)
+        print ("Results for", model_name)
+
+    results_list = load_all_results()
+    results = OrderedDict()
+
+    for k in all_k:
+        results[k] = OrderedDict()
+
+        for num_binary_search_steps in all_search_steps:
+            results[k][num_binary_search_steps] = OrderedDict()
+
+            for num_iter in all_num_iter:
+                results[k][num_binary_search_steps][num_iter] = OrderedDict()
+
+                for method_name in all_methods:
+                    filter = {
+                        "loss": eq(method_name),
+                        "model": eq(model_name),
+                        "k": eq(k),
+                        "unguided_iterations": eq(num_iter),
+                        "binary_search_steps": eq(num_binary_search_steps)
+                    }
+
+                    best_param_set, filtered_results = resolve_nonunique_filter(filter, results_list)
+
+                    if verbose and best_param_set is not None:
+                        print (f"K={k} Lr={best_param_set['unguided_lr']} and loss_coef={best_param_set['topk_loss_coef_upper']} ")
+                
+                    if best_param_set is None:
+                        continue
+
+                    assert len(filtered_results) == 5
+                    
+                    combined_results = get_combined_results(filtered_results)
+
+                    for key in list(combined_results):
+                        if "L1" not in key and "L2" not in key and "L_inf" not in key and "ASR" not in key:
+                            del combined_results[key]
+
+                    for key in list(combined_results):
+                        if "mean" not in key and "worst" not in key and "best" not in key:
+                            del combined_results[key]
+
+                    results[k][num_binary_search_steps][num_iter][method_name] = combined_results
+
+    return results
+
+if __name__ == "__main__":
+    build_full_results_dict(model_name="resnet50", verbose=True, all_search_steps=[1], all_methods=["cvxproj"])
+    build_full_results_dict(model_name="densenet121", verbose=True, all_search_steps=[1], all_methods=["cvxproj"])
+    build_full_results_dict(model_name="deit_small", verbose=True, all_search_steps=[1], all_methods=["cvxproj"])
+    build_full_results_dict(model_name="vit_base", verbose=True, all_search_steps=[1], all_methods=["cvxproj"])
+    x = 5

modelguidedattacks/run.py

@@ -0,0 +1,140 @@
+import sys
+from pprint import pformat
+from typing import Any
+import os
+import torch
+
+import ignite.distributed as idist
+import yaml
+from ignite.engine import Events
+from ignite.metrics import Accuracy, Loss
+from ignite.utils import manual_seed
+from torch import nn, optim
+
+from modelguidedattacks.data.setup import setup_data
+from modelguidedattacks.losses.boilerplate import BoilerplateLoss
+from modelguidedattacks.losses.energy import Energy, EnergyLoss
+from modelguidedattacks.metrics.topk_accuracy import TopKAccuracy
+from modelguidedattacks.models import setup_model
+from modelguidedattacks.trainers import setup_evaluator, setup_trainer
+from modelguidedattacks.utils import setup_parser, setup_output_dir
+from modelguidedattacks.utils import setup_logging, log_metrics, Engine
+
+def run(local_rank: int, config: Any):
+
+    print ("Running ", local_rank)
+    # make a certain seed
+    rank = idist.get_rank()
+    manual_seed(config.seed + rank)
+
+    # create output folder
+    config.output_dir = setup_output_dir(config, rank)
+
+    # setup engines logger with python logging
+    # print training configurations
+    logger = setup_logging(config)
+    logger.info("Configuration: \n%s", pformat(vars(config)))
+    (config.output_dir / "config-lock.yaml").write_text(yaml.dump(config))
+
+    # donwload datasets and create dataloaders
+    dataloader_train, dataloader_eval = setup_data(config, rank)
+
+    # model, optimizer, loss function, device
+    device = idist.device()
+    model = idist.auto_model(setup_model(config, idist.device()))
+    loss_fn = BoilerplateLoss().to(device=device)
+    l2_energy_loss = Energy(p=2).to(device)
+    l1_energy_loss = Energy(p=1).to(device)
+    l_inf_energy_loss = Energy(p=torch.inf).to(device)
+
+    evaluator = setup_evaluator(config, model, device)
+    evaluator.logger = logger
+
+    # attach metrics to evaluator
+    accuracy = TopKAccuracy(device=device)
+    metrics = {
+        "ASR": accuracy,
+        "L2 Energy": EnergyLoss(l2_energy_loss, device=device),
+        "L1 Energy": EnergyLoss(l1_energy_loss, device=device),
+        "L_inf Energy": EnergyLoss(l_inf_energy_loss, device=device),
+
+        "L2 Energy Min": EnergyLoss(l2_energy_loss, reduction="min", device=device),
+        "L1 Energy Min": EnergyLoss(l1_energy_loss, reduction="min", device=device),
+        "L_inf Energy Min": EnergyLoss(l_inf_energy_loss, reduction="min", device=device),
+
+        "L2 Energy Max": EnergyLoss(l2_energy_loss, reduction="max", device=device),
+        "L1 Energy Max": EnergyLoss(l1_energy_loss, reduction="max", device=device),
+        "L_inf Energy Max": EnergyLoss(l_inf_energy_loss, reduction="max", device=device)
+    }
+    for name, metric in metrics.items():
+        metric.attach(evaluator, name)
+
+    if config.guide_model in ["unguided", "instance_guided"]:
+
+        first_batch_passed = False
+        early_stopped = False
+
+        def compute_metrics(engine: Engine, tag: str):
+            nonlocal first_batch_passed
+            nonlocal early_stopped
+
+            for name, metric in metrics.items():
+                metric.completed(engine, name)
+
+            if not first_batch_passed:
+                if engine.state.metrics["ASR"] < 1e-3:
+                    print ("Early stop, assuming no success throughout")
+                    early_stopped = True
+                    engine.terminate()
+                else:
+                    first_batch_passed = True
+
+        evaluator.add_event_handler(
+            Events.ITERATION_COMPLETED(every=config.log_every_iters),
+            compute_metrics,
+            tag="eval",
+        )
+
+        evaluator.add_event_handler(
+            Events.ITERATION_COMPLETED(every=config.log_every_iters),
+            log_metrics,
+            tag="eval",
+        )
+
+        evaluator.run(dataloader_eval, epoch_length=config.eval_epoch_length)
+        log_metrics(evaluator, "eval")
+
+        if len(config.out_dir) > 0:
+            # Store results in out_dir
+            os.makedirs(config.out_dir, exist_ok=True)
+            metrics_dict = evaluator.state.metrics
+            metrics_dict["config"] = config
+            metrics_dict["early_stopped"] = early_stopped
+
+            metrics_file_path = os.path.join(config.out_dir, "results.save")
+            torch.save(metrics_dict, metrics_file_path)
+
+        # No need to train with an unguided model
+        return
+
+    assert False, "This code path is for the future"
+
+# main entrypoint
+def launch(config=None):
+    if config is None:
+        config_path = sys.argv[1]
+        config = setup_parser(config_path).parse_args(sys.argv[2:])
+        
+    backend = config.backend
+    nproc_per_node = config.nproc_per_node
+
+    if nproc_per_node == 0 or backend is None:
+        backend = None
+        nproc_per_node = None
+
+    with idist.Parallel(backend, nproc_per_node) as p:
+        p.run(run, config=config)
+
+
+if __name__ == "__main__":
+    launch()

modelguidedattacks/trainers.py

@@ -0,0 +1,83 @@
+from typing import Any, Union
+
+import ignite.distributed as idist
+import torch
+from ignite.engine import DeterministicEngine, Engine, Events
+from torch.cuda.amp import autocast
+from torch.nn import Module
+from torch.optim import Optimizer
+from torch.utils.data import DistributedSampler, Sampler
+
+
+def setup_trainer(
+    config: Any,
+    model: Module,
+    optimizer: Optimizer,
+    loss_fn: Module,
+    device: Union[str, torch.device],
+    train_sampler: Sampler,
+) -> Union[Engine, DeterministicEngine]:
+    def train_function(engine: Union[Engine, DeterministicEngine], batch: Any):
+        if config.overfit:
+            # No batch norm
+            model.eval()
+        else:
+            model.train()
+
+        samples = batch[0].to(device, non_blocking=True)
+        targets = batch[1].to(device, non_blocking=True)
+        attack_targets = batch[2].to(device, non_blocking=True)
+        sample_ids = batch[3].to(device, non_blocking=True)
+
+        with autocast(config.use_amp):
+            outputs = model(samples, attack_targets)
+            loss = loss_fn(outputs, attack_targets, targets)
+
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad()
+
+        train_loss = loss.item()
+        engine.state.metrics = {
+            "epoch": engine.state.epoch,
+            "train_loss": train_loss,
+        }
+        return {"train_loss": train_loss}
+
+
+    trainer = Engine(train_function)
+
+    # set epoch for distributed sa5mpler
+    @trainer.on(Events.EPOCH_STARTED)
+    def set_epoch():
+        if idist.get_world_size() > 1 and isinstance(train_sampler, DistributedSampler):
+            train_sampler.set_epoch(trainer.state.epoch - 1)
+
+    return trainer
+
+
+def setup_evaluator(
+    config: Any,
+    model: Module,
+    device: Union[str, torch.device],
+) -> Engine:
+    @torch.no_grad()
+    def eval_function(engine: Engine, batch: Any):
+        model.eval()
+
+        samples, gt_labels, attack_targets, sample_ids = batch
+
+        samples = samples.to(device, non_blocking=True)
+        gt_labels = gt_labels.to(device, non_blocking=True)
+        attack_targets = attack_targets.to(device, non_blocking=True)
+        sample_ids = sample_ids.to(device, non_blocking=True)
+
+        with autocast(config.use_amp):
+            outputs, perturbations = model(samples, attack_targets, gt_labels)
+
+        return outputs, attack_targets, {
+            "gt_targets": gt_labels,
+            "perturbations": perturbations
+        }
+
+    return Engine(eval_function)

modelguidedattacks/utils.py

@@ -0,0 +1,133 @@
+import logging
+from argparse import ArgumentParser
+from datetime import datetime
+from logging import Logger
+from pathlib import Path
+from typing import Any, Mapping, Optional, Union
+
+import ignite.distributed as idist
+import torch
+import yaml
+from ignite.contrib.engines import common
+from ignite.engine import Engine
+from ignite.engine.events import Events
+from ignite.handlers import Checkpoint, DiskSaver, global_step_from_engine
+from ignite.handlers.early_stopping import EarlyStopping
+from ignite.handlers.terminate_on_nan import TerminateOnNan
+from ignite.handlers.time_limit import TimeLimit
+from ignite.utils import setup_logger
+
+
+def setup_parser(config_path="base_config.yaml"):
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f.read())
+
+    parser = ArgumentParser()
+    parser.add_argument("--config", default=None, type=str)
+    parser.add_argument("--backend", default=None, type=str)
+    for k, v in config.items():
+        if isinstance(v, bool):
+            parser.add_argument(f"--{k}", action="store_true")
+        else:
+            parser.add_argument(f"--{k}", default=v, type=type(v))
+
+    return parser
+
+
+def log_metrics(engine: Engine, tag: str) -> None:
+    """Log `engine.state.metrics` with given `engine` and `tag`.
+
+    Parameters
+    ----------
+    engine
+        instance of `Engine` which metrics to log.
+    tag
+        a string to add at the start of output.
+    """
+    metrics_format = "{0} [{1}/{2}]: {3}".format(
+        tag, engine.state.epoch, engine.state.iteration, engine.state.metrics
+    )
+
+    epoch_size = engine.state.epoch_length
+    local_iteration = engine.state.iteration - epoch_size * (engine.state.epoch - 1)
+    metrics_format = f"{tag} Epoch {engine.state.epoch} - [{local_iteration} / {epoch_size}] : {engine.state.metrics}"
+
+    engine.logger.info(metrics_format)
+
+
+def resume_from(
+    to_load: Mapping,
+    checkpoint_fp: Union[str, Path],
+    logger: Logger,
+    strict: bool = True,
+    model_dir: Optional[str] = None,
+) -> None:
+    """Loads state dict from a checkpoint file to resume the training.
+
+    Parameters
+    ----------
+    to_load
+        a dictionary with objects, e.g. {“model”: model, “optimizer”: optimizer, ...}
+    checkpoint_fp
+        path to the checkpoint file
+    logger
+        to log info about resuming from a checkpoint
+    strict
+        whether to strictly enforce that the keys in `state_dict` match the keys
+        returned by this module’s `state_dict()` function. Default: True
+    model_dir
+        directory in which to save the object
+    """
+    if isinstance(checkpoint_fp, str) and checkpoint_fp.startswith("https://"):
+        checkpoint = torch.hub.load_state_dict_from_url(
+            checkpoint_fp,
+            model_dir=model_dir,
+            map_location="cpu",
+            check_hash=True,
+        )
+    else:
+        if isinstance(checkpoint_fp, str):
+            checkpoint_fp = Path(checkpoint_fp)
+
+        if not checkpoint_fp.exists():
+            raise FileNotFoundError(f"Given {str(checkpoint_fp)} does not exist.")
+        checkpoint = torch.load(checkpoint_fp, map_location="cpu")
+
+    Checkpoint.load_objects(to_load=to_load, checkpoint=checkpoint, strict=strict)
+    logger.info("Successfully resumed from a checkpoint: %s", checkpoint_fp)
+
+
+def setup_output_dir(config: Any, rank: int) -> Path:
+    """Create output folder."""
+    if rank == 0:
+        now = datetime.now().strftime("%Y%m%d-%H%M%S")
+        name = f"{now}-backend-{config.backend}-lr-{config.lr}"
+        path = Path(config.output_dir, name)
+        path.mkdir(parents=True, exist_ok=True)
+        config.output_dir = path.as_posix()
+
+    return Path(idist.broadcast(config.output_dir, src=0))
+
+
+def setup_logging(config: Any) -> Logger:
+    """Setup logger with `ignite.utils.setup_logger()`.
+
+    Parameters
+    ----------
+    config
+        config object. config has to contain `verbose` and `output_dir` attribute.
+
+    Returns
+    -------
+    logger
+        an instance of `Logger`
+    """
+    green = "\033[32m"
+    reset = "\033[0m"
+    logger = setup_logger(
+        name=f"{green}[ignite]{reset}",
+        level=logging.DEBUG if config.debug else logging.INFO,
+        format="%(name)s: %(message)s",
+        filepath=config.output_dir / "training-info.log",
+    )
+    return logger

print_results.py

@@ -0,0 +1,25 @@
+from modelguidedattacks import results
+
+results_list = results.load_all_results()
+
+filter = {
+    "loss": results.in_set(["cwk"]),
+    "model": results.eq("vit_base"),
+    "k": results.eq(5),
+    "binary_search_steps": results.eq(1),
+    "unguided_iterations": results.eq(30),
+    # "topk_loss_coef_upper": results.eq(20),
+    # "unguided_lr": results.eq(0.002),
+    "cvx_proj_margin": results.eq(0.2),
+    # "seed": results.eq(10),
+}
+
+filtered_results = results.filter_results(filter, results_list)
+
+print ("Found", len(filtered_results))
+
+for result in filtered_results:
+    print ("-" * 30)
+    for key, val in result.items():
+        print (key, "=", val)
+    print ("-" * 30)

print_table.py

@@ -0,0 +1,163 @@
+from pylatex import Document, Section, Subsection, Tabular, MultiColumn,\
+    MultiRow, NoEscape
+from pylatex.math import Math
+from collections import OrderedDict
+import copy
+import numpy as np
+
+from modelguidedattacks.results import build_full_results_dict
+
+def result_to_str(result, long=False):
+    if result is None or np.isinf(result) or np.isnan(result):
+        return "-"
+    elif long:
+        return f"{result:.4f}"
+    else:
+        return f"{result:.2f}"
+
+"""
+results top level will be keyed by K
+next level will be binary search steps
+next level will be keyed by iterations
+next level will be keyed by method
+"""
+
+only_mean = False
+model_name = "resnet50"
+results = build_full_results_dict(model_name)
+
+model_to_tex = {
+    "resnet50": "Resnet-50",
+    "densenet121": "Densenet121",
+    "deit_small": "DeiT-S",
+    "vit_base": "ViT$_{B}$"
+}
+
+# Preprocess all results and select bests
+for top_k, bs_dict in results.items():
+    for num_bs, iter_dict in bs_dict.items():
+        for num_iter, method_dict in iter_dict.items():
+            metric_bests = {}
+            metrics_compared = {}
+
+            for method_name, method_results in method_dict.items():
+                for metric_name, metric_value in method_results.items():
+                    reduction_func = max if "ASR" in metric_name else min
+
+                    if metric_name not in metric_bests:
+                        metric_bests[metric_name] = 0. if reduction_func is max else np.Infinity
+                        metrics_compared[metric_name] = 0
+
+                    if metric_value is not None:
+                        metric_bests[metric_name] = reduction_func(metric_bests[metric_name], metric_value)
+                        metrics_compared[metric_name] += 1
+
+            for method_name, method_results in method_dict.items():
+                for metric_name, metric_value in method_results.items():
+                    method_results[metric_name] = result_to_str(metric_value, "inf" in metric_name or "ASR" in metric_name)
+
+                    if metric_value is not None and np.allclose(metric_value, metric_bests[metric_name]) \
+                        and metrics_compared[metric_name] > 1:
+                        method_results[metric_name] = rf"\textbf{{ {method_results[metric_name]} }}"
+
+method_tex = {
+    "cwk": r"CW^K",
+    "ad": r"AD",
+    "cvxproj": r"\textbf{QuadAttac$K$}"
+}
+
+doc = Document("multirow")
+
+protocol_cols = 1
+attack_method_cols = 1
+best_cols = 4
+mean_cols = 4
+worst_cols = 4
+
+if only_mean:
+    col_widths = [protocol_cols, attack_method_cols, mean_cols]
+else:
+    col_widths = [protocol_cols, attack_method_cols, best_cols, mean_cols, worst_cols]
+
+total_cols = sum(col_widths)
+tabular_string = "|"
+
+for w in col_widths:
+    tabular_string += "l" * w + "|"
+
+table1 = Tabular(tabular_string)
+table1.add_hline()
+table1.add_row((MultiColumn(total_cols, align='|c|', data=NoEscape(model_to_tex[model_name])),))
+table1.add_hline()
+
+if only_mean:
+    table1.add_row((
+        MultiRow(2, data="Protocol"),
+        MultiRow(2, data="Attack Method"),
+        MultiColumn(mean_cols, align="|c|", data="Mean"),
+    ))
+else:
+    table1.add_row((
+        MultiRow(2, data="Protocol"),
+        MultiRow(2, data="Attack Method"),
+        MultiColumn(best_cols, align="|c|", data="Best"),
+        MultiColumn(mean_cols, align="|c|", data="Mean"),
+        MultiColumn(worst_cols, align="|c|", data="Worst"),
+    ))
+
+table1.add_hline(start=protocol_cols + attack_method_cols + 1)
+
+num_result_colums = 1 if only_mean else 3
+table1.add_row("", 
+               "",
+               *(NoEscape(r"ASR$\uparrow$"),
+               NoEscape(r"$\ell_1 \downarrow$"),
+               NoEscape(r"$\ell_2 \downarrow$"),
+               NoEscape(r"$\ell_{\infty} \downarrow$"))*num_result_colums
+               )
+
+table1.add_hline()
+
+for top_k, bs_dict in results.items():
+    
+    total_results = 0
+    # Count total results
+    for _, iter_dict in bs_dict.items():
+        for num_iter, method_dict in iter_dict.items():
+            total_results += len(method_dict)
+
+    top_k_latex_obj = MultiRow(total_results, data=f"Top-{top_k}")
+
+    shown_topk_obj = False
+
+    for bs_steps, iter_dict in bs_dict.items():
+        for num_iter, method_dict in iter_dict.items():
+            for method_name, method_results in method_dict.items():
+                first_obj = top_k_latex_obj if not shown_topk_obj else ""
+                shown_topk_obj = True
+
+                row_results = []
+
+                reduction_names = ["mean"] if only_mean else ["best", "mean", "worst"]
+                for reduction in reduction_names:
+                    for metric in ["ASR", "L1", "L2", "L_inf"]:
+                        result_key = f"{metric}_{reduction}"
+                        long_result = "inf" in metric
+                        row_results.append(
+                            NoEscape(method_results[result_key])
+                            )
+
+                table1.add_row(
+                    first_obj,
+                    NoEscape("$" + method_tex[method_name] +
+                            f"_{{{bs_steps}x{num_iter}}}$"),
+                    *row_results
+                )
+
+            table1.add_hline(start=protocol_cols + 1)
+
+    table1.add_hline()
+
+# doc.append(table1)
+
+print(table1.dumps())

result_stats.py

@@ -0,0 +1,90 @@
+import numpy as np
+import sys
+from modelguidedattacks import results
+
+results_list = results.load_all_results()
+
+filter = {
+    "loss": results.in_set(["cvxproj"]),
+    "model": results.eq("resnet50"),
+    "k": results.eq(20),
+    "binary_search_steps": results.eq(1),
+    "unguided_iterations": results.eq(60),
+    # "topk_loss_coef_upper": results.eq(20),
+    # "unguided_lr": results.eq(0.002),
+    "cvx_proj_margin": results.eq(0.2),
+    "topk_loss_coef_upper": results.gte(12)
+    # "seed": results.eq(10),
+}
+
+filtered_results = results.filter_results(filter, results_list)
+
+combined_results = {}
+
+for result in filtered_results:
+    for key in result:
+        if key not in combined_results:
+            combined_results[key] = []
+
+        combined_results[key].append(result[key])
+
+unique_runs = len(np.unique(combined_results["seed"]))
+print ("Stats from", len(filtered_results))
+# assert len(combined_results["seed"]) == unique_runs
+
+for key, val in list(combined_results.items()):
+    if key in ["ASR", "L1", "L2", "L_inf"]:
+        val = np.array(val)
+        combined_results[f"{key}_mean"] = np.mean(val[np.isfinite(val)])
+        combined_results[f"{key}_median"] = np.median(val[np.isfinite(val)])
+
+# Coupled results
+best_asr_idx = np.argmax(combined_results["ASR"])
+best_asr = combined_results["ASR"][best_asr_idx]
+best_l1 = combined_results["L1"][best_asr_idx]
+best_l2 = combined_results["L2"][best_asr_idx]
+best_linf = combined_results["L_inf"][best_asr_idx]
+
+combined_results["ASR_best"] = best_asr
+combined_results["L1_best"] = best_l1
+combined_results["L2_best"] = best_l2
+combined_results["L_inf_best"] = best_linf
+
+worst_asr_idx = np.argmin(combined_results["ASR"])
+worst_asr = combined_results["ASR"][worst_asr_idx]
+worst_l1 = combined_results["L1"][worst_asr_idx]
+worst_l2 = combined_results["L2"][worst_asr_idx]
+worst_linf = combined_results["L_inf"][worst_asr_idx]
+
+combined_results["ASR_worst"] = worst_asr
+combined_results["L1_worst"] = worst_l1
+combined_results["L2_worst"] = worst_l2
+combined_results["L_inf_worst"] = worst_linf
+
+draw_keys = ["best", "mean", "worst"]
+val_keys = ["ASR", "L1", "L2", "L_inf"]
+
+print ("---------------")
+for draw_key in draw_keys:
+    for val_key in val_keys:
+        key = val_key + "_" + draw_key
+        val = combined_results[key]
+
+        print (key, val)
+print ("---------------")
+
+for draw_key in draw_keys:
+    for val_key in val_keys:
+        key = val_key + "_" + draw_key
+        val = combined_results[key]
+
+        if np.isinf(val):
+            val = "N/A"
+
+        sep = "&"
+        if isinstance(val, str):
+            sys.stdout.write(f"{val} {sep} ")    
+        elif "inf" in key:
+            sys.stdout.write(f"{val:.3f} {sep} ")
+        else:
+            sys.stdout.write(f"{val:.2f} {sep} ")

setup.py

@@ -0,0 +1,40 @@
+import io
+import os
+from setuptools import find_packages, setup
+
+def read(*paths, **kwargs):
+    """Read the contents of a text file safely.
+    >>> read("project_name", "VERSION")
+    '0.1.0'
+    >>> read("README.md")
+    ...
+    """
+
+    content = ""
+    with io.open(
+        os.path.join(os.path.dirname(__file__), *paths),
+        encoding=kwargs.get("encoding", "utf8"),
+    ) as open_file:
+        content = open_file.read().strip()
+    return content
+
+
+def read_requirements(path):
+    return [
+        line.strip()
+        for line in read(path).split("\n")
+        if not line.startswith(('"', "#", "-", "git+"))
+    ]
+
+
+setup(
+    name="modelguidedattacks",
+    version="0.1",
+    description="Adversarial attacks",
+    url="",
+    long_description=read("README.md"),
+    long_description_content_type="text/markdown",
+    author="NCSU",
+    packages=["modelguidedattacks"],
+    install_requires=read_requirements("requirements.txt"),
+)