Upload 192 files

README.md

@@ -1,3 +1,113 @@
----
-license: apache-2.0
----
+# SAE: Sustainable Adversarial Example Evaluation Framework for Class-Incremental Learning
+
+## 🌟 Overview
+
+_**News**: This work has been accepted as a poster by [AAAI 2026](https://aaai.org/conference/aaai/aaai-26/)._
+
+**SAE (Sustainable Adversarial Example)** is a *universal adversarial attack framework* targeting **Class-Incremental Learning (CIL)**. This repository provides a comprehensive pipeline for both CIL training and benchmarking multiple attack methods, including our proposed SAE approach.
+
+The project integrates with [PyCIL: A Python Toolbox for Class-Incremental Learning](https://github.com/LAMDA-CL/PyCIL) for CIL model training. It also supports benchmarking several attack baselines alongside SAE, enabling fair and reproducible evaluations of adversarial robustness across CIL methods.
+
+If you are interested in our work, please refer to:
+```
+@inproceedings{liu2026SAE,
+  title={Improving Sustainability of Adversarial Examples in Class-Incremental Learning},
+  author={Taifeng Liu, Xinjing Liu, Liangqiu Dong, Yang Liu, Yilong Yang, Zhuo Ma},
+  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
+  year={2026}
+}
+```
+
+---
+
+## ⚙️ Environment Setup
+
+### Project Layout
+
+* `attacks/`: Implementations of all attack baselines including **MIFGSM**, **Gaker**, **AIM**, **CGNC**, **CleanSheet**, **UnivIntruder**, and **SAE**.
+* `convs/`: Backbone definitions for CIL models and CLIP model, including `resnet32`, `resnet50`, `cosine_resnet32`, and `cosine_resnet50`.
+  * **Pre-trained CLIP** is downloaded from [HuggingFace](https://huggingface.co/laion/CLIP-ViT-B-32-laion2B-s34B-b79K).
+* `datasets/`: Dataset management for CIFAR-100 (32x32) and ImageNet-100 (224x224).
+  * **CIFAR-100** is automatically downloaded at runtime.
+  * **ImageNet-100** should be manually extracted from ImageNet-1K using `create_imagenet100_from_imagenet.py`, which parses `train.txt` and `eval.txt` to extract relevant classes.
+* `exps/`: JSON configuration files for various CIL training methods.
+* `logs/`: Stores trained CIL model checkpoints and evaluation results of attacks.
+* `models/`: Contains implementations of **9** CIL algorithms: **BiC**, **DER**, **Finetune**, **Foster**, **iCaRL**, **MEMO**, **PodNet**, **Replay**, **WA**.
+* `scripts/`: Shell scripts for CIL training and attack benchmarking.
+* `utils/`: Utility functions for augmentation, logging, dataset processing, visualization, etc.
+* `attack.py`: Main entry point to run adversarial attacks.
+* `trainCIL.py`: Main entry point to train CIL models.
+
+### Dependency Requirements
+
+* **OS**: Ubuntu 22.04
+* **Python**: 3.12
+* **PyTorch**: ≥ 2.1
+* **GPU**: NVIDIA RTX 4090 (24GB VRAM)
+
+To set up the environment:
+
+```bash
+conda env create -f environment.yml
+conda activate SAE
+```
+
+---
+
+## 🚀 Result Reproduction
+
+### Step 1: CIL Training
+
+First, you need to train the target CIL models follow the instruction below.
+
+Example: training a single CIL method (e.g., **iCaRL**) on **CIFAR-100**:
+
+```bash
+python trainCIL.py --config exps/icarl.json
+```
+
+Example: training a single CIL method (e.g., **iCaRL**) on **ImageNet-100**:
+
+```bash
+python trainCIL.py --config exps/icarl-imagenet100.json
+```
+
+We also provide scripts to train all 9 CIL methods on **CIFAR-100**:
+
+```bash
+./scripts/trainCIL-CIFAR100.sh
+```
+
+Train all 9 CIL methods on **ImageNet-100**:
+
+```bash
+./scripts/trainCIL-ImageNet100.sh
+```
+
+All model checkpoints will be saved under the `logs/` directory, organized by method and dataset.
+
+
+---
+
+### Step 2: Adversarial Attack Benchmarking
+
+_Note: assume that you have already prepared the dataset, the CIL model, and the CLIP model._
+
+To launch an **SAE** attack targeting class **0** on a CIL model trained on **CIFAR-100**:
+
+```bash
+python attack.py --config exps/icarl.json --attack_method SAE --target_class 0
+```
+
+To test a different attack baseline, simply change the `--attack_method` argument.
+
+You can reproduce the overall evaluation results by runing the below scripts (referring to the **Table 1** in our paper).
+
+```bash
+./scripts/attacks-CIFAR100.sh
+./scripts/attacks-ImageNet100.sh
+```
+
+#### Benchmark Results
+
+All benchmark results are available in the `appendix.pdf` of Supplementary Material.

attack.py

@@ -0,0 +1,140 @@
+import json
+import argparse
+import sys
+import logging
+import copy
+import torch
+import os
+
+from trainCIL import _set_random, _set_device, print_args
+from attacks.AIM.AIMAttack import AIM
+from attacks.BASEAttack import BASEAttack
+from attacks.Gaker.GAKERAttack import Gaker
+from attacks.CGNC.CGNCAttack import CGNC
+from attacks.CleanSheet.CleanSheetAttack import CleanSheet
+from attacks.UnivIntruder.UnivIntruderAttack import UnivIntruder
+from attacks.SAE.SAEAttack import SAE
+
+
+def main():
+    args = setup_parser().parse_args()
+    param = load_json(args.config)
+    args = vars(args)  # Converting argparse Namespace to a dict.
+    args.update(param)  # Add parameters from json
+
+    evaluate(args)
+
+
+def evaluate(args):
+    seed_list = copy.deepcopy(args["seed"])
+    device = copy.deepcopy(args["device"])
+
+    for seed in seed_list:
+        args["seed"] = seed
+        args["device"] = device
+        _evaluate(args)
+
+
+def _evaluate(args):
+    # For attacks
+    args["attack"] = True
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    init_cls = 0 if args["init_cls"] == args["increment"] else args["init_cls"]
+    logs_name = "logs/{}/{}/init_cls_{}/per_classes_{}/{}".format(args["model_name"], args["dataset"], init_cls,
+                                                                         args['increment'], args["convnet_type"])
+    if not os.path.exists(logs_name):
+        raise "Model is not trained yet."
+    logs_eval_name = "logs/{}/{}/init_cls_{}/per_classes_{}/{}/eval/{}".format(args["model_name"], args["dataset"], init_cls,
+                                                               args['increment'], args["convnet_type"], args['attack_method'])
+    if not os.path.exists(logs_eval_name):
+        os.makedirs(logs_eval_name)
+
+    logfilename = "logs/{}/{}/init_cls_{}/per_classes_{}/{}/eval/{}/adv_log".format(
+        args["model_name"],
+        args["dataset"],
+        init_cls,
+        args["increment"],
+        args["convnet_type"],
+        args['attack_method']
+    )
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(filename)s] => %(message)s",
+        handlers=[
+            logging.FileHandler(filename=logfilename + ".log"),
+            logging.StreamHandler(sys.stdout),
+        ],
+    )
+
+    args['epsilons'] = [0.01, 0.015, 0.03, 0.06, 0.1, 0.2]
+    args['logs_name'] = logs_name
+    args['logs_eval_name'] = logs_eval_name
+    _set_random()
+    _set_device(args)
+    print_args(args)
+
+    # Init the attack
+    if args['attack_method'] in NEW_ATTACKS:
+        adv = init_new_attack(args, device=device)
+        if args['attack_method'] == 'AIM' or args['attack_method'] == 'Gaker' or args['attack_method'] == 'CGNC':
+            adv.train_generator()
+        elif args['attack_method'] == 'SAE' or args['attack_method'] == 'UnivIntruder' or args['attack_method'] == 'CleanSheet':
+            adv.train_adv()
+    else:
+        adv = init_foolbox_attack(args, device=device)
+
+    # Conduct attack
+    adv.run_test()
+
+
+NEW_ATTACKS = ['AIM', 'Gaker', 'CGNC', 'CleanSheet', 'UnivIntruder', 'SAE']
+FOOLBOX_ATTACKS = ['L2FGM', 'FGSM', 'MIFGSM',
+                   'L1PGD', 'L2PGD', 'LinfPGD',
+                   'L2DeepFool', 'LinfDeepFool', 'BoundaryAttack',
+                   'CarliniWagnerL2', 'GaussianNoise', 'UniformNoise']
+
+def init_new_attack(args, device='cuda', **kwargs):
+    attack_name = args['attack_method']
+    models = kwargs.get('models', None)
+
+    if attack_name == 'AIM':
+        adv = AIM(args=args, device=device)
+    elif attack_name == 'Gaker':
+        adv = Gaker(args=args, device=device)
+    elif attack_name == 'CGNC':
+        adv = CGNC(args=args, device=device)
+    elif attack_name == 'CleanSheet':
+        adv = CleanSheet(args=args, device=device)
+    elif attack_name == 'UnivIntruder':
+        adv = UnivIntruder(args=args, device=device)
+    elif attack_name == 'SAE':
+        adv = SAE(args=args, device=device)
+    else:
+        raise ValueError(f"Unknown attack method: {attack_name}")
+
+    return adv
+
+def init_foolbox_attack(args, device='cuda', **kwargs):
+    attack = BASEAttack(args=args, device=device)
+    return attack
+
+def load_json(settings_path):
+    with open(settings_path) as data_file:
+        param = json.load(data_file)
+
+    return param
+
+
+def setup_parser():
+    parser = argparse.ArgumentParser(description='Reproduce of multiple continual learning algorithms.')
+    parser.add_argument('--config', type=str, default='exps/finetune.json', help='Json file of settings.')
+    parser.add_argument('--batch_size', type=int, default=128, help='set the batch size.')
+    parser.add_argument('--attack_method', type=str, default='AIM', help='set the attack method, e.g., LinfPGD, MIFGSM, AIM, Gaker, CGNC, CleanSheet, UnivIntruder, SAE.')
+    parser.add_argument('--target_class', type=int, default=0, help='the target class, None indicates untargeted attack.')
+    parser.add_argument('--eval', action='store_true', help='evaluation only')
+    return parser
+
+
+if __name__ == '__main__':
+    main()

attacks/AIM/AIMAttack.py

@@ -0,0 +1,226 @@
+import os.path
+import torch
+from foolbox.attacks.base import *
+from foolbox.attacks.gradient_descent_base import *
+from tqdm import tqdm
+from attacks.AIM.src.gat.models.attack import AIMAttack, ContrastiveLoss
+from attacks.AIM.src.gat.models.surrogate import midlayer_dict, register_collecter, register_collecter_cl
+from attacks.attack_config import SustainableAttack
+from utils.plot import plot_asr_per_target, save_grad_cam
+import logging
+import pandas as pd
+import foolbox as fb
+from foolbox import PyTorchModel
+import numpy as np
+from utils import factory
+from utils.data_manager import get_dataloader
+
+
+class AIM(SustainableAttack):
+    def __init__(self, args, device='cuda'):
+        super().__init__(args, device)
+        self.device = device
+        self.args = args
+        self.surrogate_model = None
+
+        self.adv_generator = AIMAttack(device=device)
+        self.adv_generator.set_mode('train')
+        self.lr = 0.001
+        self.betas = (0.5, 0.999)
+        self.num_epoch = 100
+        self.optim = torch.optim.Adam(self.adv_generator.get_params(), lr=self.lr, betas=self.betas)
+        self.contrastive_loss = ContrastiveLoss(0.2)
+        self.sim_loss = torch.nn.functional.cosine_similarity
+        self.eval_batch_szie = 128
+
+        self.surrogate_model_name = 'resnet32_cl'
+        self.layer = midlayer_dict[self.surrogate_model_name]
+        self.prefix = (f'adv_generator_{self.surrogate_model_name}'
+                       f'_{self.layer}'
+                       f'_tclass{self.target_class}')
+        self.save_path = os.path.join(self.args['logs_eval_name'], f'target{str(self.target_class)}')
+        if not os.path.exists(self.save_path):
+            os.makedirs(self.save_path)
+
+        self.plot_gradcam = False
+
+    def train_generator(self):
+        if 'cl' in self.surrogate_model_name:
+            s_model = factory.get_model(self.args["model_name"], self.args)
+            s_model.incremental_train(self.data_manager)
+            s_model._network.load_state_dict(
+                torch.load(self.ckpt_paths[0], map_location=self.device)['model_state_dict'])
+            s_model._network.to(self.device)
+            s_model._network.eval()
+            self.surrogate_model = s_model._network
+            del s_model
+            torch.cuda.empty_cache()
+            self.feat_collecter = []
+            self.feat_collecter_handler, self.feat_collecter = register_collecter_cl(self.surrogate_model,
+                                                                                     self.layer,
+                                                                                     self.feat_collecter,
+                                                                                     self.args["model_name"])
+        else:
+            self.surrogate_model = torch.hub.load("chenyaofo/pytorch-cifar-models", 'cifar100_resnet32', pretrained=True)
+            self.surrogate_model.to(self.device)
+            self.surrogate_model.eval()
+            self.feat_collecter = []
+            self.feat_collecter_handler, self.feat_collecter = register_collecter(self.surrogate_model,
+                                                                                  self.layer,
+                                                                                  self.feat_collecter)
+        self.file_path = os.path.join(self.save_path, f'{self.prefix}.pth')
+        if os.path.exists(self.file_path):
+            self.adv_generator.load_ckpt(self.file_path)
+            self.adv_generator.set_mode('eval')
+        else:
+            loaders = get_dataloader(self.data_manager, batch_size=self.batch_size,
+                                start_class=0, end_class=10,
+                                train=True, shuffle=True, num_workers=0)
+
+            target_images = []
+            target_labels = []
+            for data in loaders:
+                _, image_batch, label_batch = data
+                mask = label_batch == self.target_class
+                selected_images = image_batch[mask]
+                selected_labels = label_batch[mask]
+                target_images.append(selected_images)
+                target_labels.append(selected_labels)
+            del loaders
+            target_images = torch.cat(target_images, dim=0).to(self.device)
+            target_labels = torch.cat(target_labels, dim=0).to(self.device)
+            target_images, target_labels = ep.astensors(*(target_images[:self.batch_size], target_labels[:self.batch_size]))
+
+            total_loss = []
+            for epoch in range(1, self.num_epoch + 1):
+                laoder_tqdm = tqdm(self.loader, total=len(self.loader), desc=f'Epoch {epoch}')
+                loss_np = 0
+                for i, (_, x, y) in enumerate(laoder_tqdm):
+                    x_f = x[y != self.target_class].to(self.device)
+                    del x, y
+                    if len(x_f) > len(target_images):
+                        x_f = x_f[:len(target_images)].to(self.device)
+                    else:
+                        random_indices = torch.randperm(len(target_images))[:len(x_f)].to(self.device)
+                        target_images = target_images[random_indices]
+
+                    x_adv = self.adv_generator(x_f, target_images.raw.to(self.device))
+
+                    logits_nat = self.surrogate_model(self.norm(x_f))
+                    feat_nat = self.feat_collecter.pop()
+                    logits_tar = self.surrogate_model(self.norm(target_images.raw))
+                    feat_tar = self.feat_collecter.pop()
+                    logits_adv = self.surrogate_model(self.norm(x_adv))
+                    feat_adv = self.feat_collecter.pop()
+
+                    loss = (self.contrastive_loss(logits_adv, logits_nat, logits_tar) +
+                            self.sim_loss(feat_nat, feat_adv) -
+                            self.sim_loss(feat_tar, feat_adv)).mean()
+                    # print(loss.item())
+                    loss_np = loss_np + loss.item()
+
+                    self.optim.zero_grad()
+                    loss.backward()
+                    self.optim.step()
+                    del x_f, x_adv, logits_nat, logits_adv, logits_tar, feat_nat, feat_tar, feat_adv
+                    torch.cuda.empty_cache()
+                total_loss.append(loss_np/(i+1))
+                logging.info(f'Epoch {epoch} loss: {loss_np/(len(self.loader))}')
+            logging.info(f'Total loss: {total_loss}')
+            self.feat_collecter_handler.remove()
+            self.adv_generator.save_ckpt(self.file_path)
+
+
+    def run_test(self):
+        # Load Batch Data
+        self.adv_generator.set_mode('eval')
+        self.adv_generator.adv_gen.to(self.device)
+        self.loader = get_dataloader(self.data_manager, batch_size=self.eval_batch_szie,
+                                start_class=0, end_class=10,
+                                train=False, shuffle=False, num_workers=0)
+        target_images = []
+        target_labels = []
+        for data in self.loader:
+            _, image_batch, label_batch = data
+            mask = label_batch == self.target_class
+            selected_images = image_batch[mask]
+            selected_labels = label_batch[mask]
+            target_images.append(selected_images)
+            target_labels.append(selected_labels)
+        target_imgs = torch.cat(target_images, dim=0).to(self.device)
+        target_labels = torch.cat(target_labels, dim=0).to(self.device)
+        target_imgs, target_labels = ep.astensors(*(target_imgs, target_labels))
+        for i, (_, imgs, labels) in enumerate(tqdm(self.loader, total=len(self.loader),
+                                                   desc=f'Loading Data with Batch Size of {self.batch_size}) :')):
+            if i > 0:
+                break
+
+            imgs, labels = ep.astensors(*(imgs.to(self.device), labels.to(self.device)))
+
+            imgs_f = imgs[labels != self.target_class]
+            labels_f = labels[labels != self.target_class]
+            labels_t_f = ep.full_like(labels_f, fill_value=self.target_class)
+
+            self.attacks(i, imgs_f, labels_f, labels_t_f, target_imgs[:20], target_labels[:20])
+
+
+    def attacks(self, i_batch, imgs, labels, labels_t, target_imgs=None, target_labels=None):
+        asr_matrix = np.ones((10, len(target_imgs)))
+        self.model = factory.get_model(self.args["model_name"], self.args)
+        for task in range(10):
+            logging.info("***** Starting attack on task [{}]. *****".format(task))
+            self.model.incremental_train(self.data_manager)
+            self.model._network.load_state_dict(torch.load(self.ckpt_paths[task], map_location=self.device)['model_state_dict'])
+            self.model._network.to(self.device)
+            self.model._network.eval()
+
+            # Run attack on ecah target image
+            criterion = fb.criteria.Misclassification(
+                labels) if self.target_class is None else fb.criteria.TargetedMisclassification(
+                labels_t)
+            current_model = PyTorchModel(self.model._network, bounds=(0, 1), preprocessing=self.preprocessing)
+            verify_input_bounds(imgs, current_model)
+            criterion = get_criterion(criterion)
+            is_adversarial = get_is_adversarial(criterion, current_model)
+
+            logging.info("Eval attack on each target images.")
+            for i, target_image in enumerate(target_imgs):
+                advs = ep.astensor(self.adv_generator(imgs.raw.to(self.device), target_image.raw.repeat(len(imgs), 1, 1, 1).to(self.device)))
+                is_adv = is_adversarial(advs)[0]
+                asr_matrix[task, i] = (is_adv.bool().sum().raw.item() / len(imgs))
+                if self.plot_gradcam:
+                    save_grad_cam(self.args, torch.clip(advs.raw.detach(),0,1), labels_t.raw,
+                                  self.model._network, self.save_path + "/GradCam" + f"targetimg{i}", prefix=f'task{task}',
+                                  layer_name='stage_3', save_num=100, save_raw=True)
+
+                del advs, is_adv, target_image
+                torch.cuda.empty_cache()
+
+            del criterion, current_model, is_adversarial
+            torch.cuda.empty_cache()
+
+            self.model.after_task()
+
+        # Save all target images info: everage asr,
+        asr_matrix = np.mean(asr_matrix, axis=1, keepdims=True)
+        prefix = f'batch{i_batch}_{self.prefix}'
+        plot_asr_per_target(asr_matrix, self.save_path, prefix, self.args)
+        df = pd.DataFrame(asr_matrix, columns=['ASR'])
+        df.to_excel(os.path.join(self.save_path, f"{prefix}.xlsx"), index=False)
+
+        del asr_matrix, imgs, labels, labels_t, target_imgs
+        torch.cuda.empty_cache()
+
+    def __call__(
+            self,
+            model: Model,
+            inputs: T,
+            criterion: Any,
+            *,
+            epsilons: Union[Sequence[Union[float, None]], float, None],
+            **kwargs: Any,
+    ) -> Union[Tuple[List[T], List[T], T], Tuple[T, T, T]]:
+        ...
+
+    def repeat(self, times: int) -> "AIM":
+        ...

attacks/AIM/examples/aim_attack.py

@@ -0,0 +1,237 @@
+#!/usr/bin/env python3
+# Usage:
+# ./examples/aim_attack.py -h
+import argparse
+import json
+from pathlib import Path
+from pprint import pformat
+from typing import List, Union
+
+import torch
+from tqdm import tqdm
+
+from attacks.GAT.src.gat.datasets import build_dataset, list_datasets
+from attacks.GAT.src.gat.datasets.transforms import norm
+from attacks.GAT.src.gat.models.attack import AIMAttack, ContrastiveLoss
+from attacks.GAT.src.gat.models.surrogate import (build_surrogate, list_surrogates,
+                                  midlayer_dict, register_collecter)
+from attacks.GAT.src.gat.runtime import AverageMeter, calc_cls_accuracy, fix_random, randid
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-v', '--verbose', action='store_true')
+    parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--expid', type=str, default=randid(4))
+    parser.add_argument('--workdir', type=str, default='workdirs')
+    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--tar-classes', type=int, default=24)
+    parser.add_argument('--batch-size', type=int, default=16)
+    parser.add_argument('--dataset',
+                        type=str,
+                        default='imagenet',
+                        choices=list_datasets())
+    parser.add_argument('--data-root',
+                        type=str,
+                        default=Path(__file__).parent / '../data' / 'in_1k')
+    sub_parsers = parser.add_subparsers(dest='command')
+    train_parser = sub_parsers.add_parser('train')
+    train_parser.add_argument('--surrogate-id',
+                              type=str,
+                              default='resnet152',
+                              choices=list_surrogates())
+    train_parser.add_argument('--num-epoch', type=int, default=10)
+    train_parser.add_argument('--lr', type=float, default=0.0002)
+    train_parser.add_argument('--betas',
+                              type=float,
+                              nargs=2,
+                              default=(0.5, 0.999))
+    sub_parsers.add_parser('evaluate')
+    args = parser.parse_args()
+
+    args.workdir = Path(args.workdir) / args.expid
+    args.workdir.mkdir(parents=True, exist_ok=True)
+    args.device = torch.device(args.device)
+
+    args.ckpt = args.workdir / 'model.pth'
+
+    fix_random(args.seed)
+
+    with open(args.workdir / 'args.txt', 'w') as f:
+        f.write(pformat(vars(args)))
+
+    return args
+
+
+def init_loader(dataset: str,
+                data_root: Union[str, Path],
+                tar_classes: Union[int, List[int]],
+                batch_size: int = 16,
+                command: str = 'train') -> List[torch.utils.data.DataLoader]:
+    train_ds = build_dataset(dataset,
+                             data_root=data_root,
+                             is_train=(command == 'train'))
+    train_loader = torch.utils.data.DataLoader(
+        train_ds,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=4,
+        pin_memory=True,
+    )
+    target_ds = build_dataset(dataset,
+                              data_root=data_root,
+                              is_train=True,
+                              filter_class=tar_classes)
+    target_loader = torch.utils.data.DataLoader(
+        target_ds,
+        batch_size=batch_size,
+        sampler=torch.utils.data.RandomSampler(target_ds,
+                                               replacement=True,
+                                               num_samples=len(train_ds)),
+        num_workers=4,
+        pin_memory=True,
+    )
+    return train_loader, target_loader
+
+
+def train(
+    surrogate_id: str,
+    dataset: str,
+    data_root: Union[str, Path],
+    tar_classes: Union[int, List[int]] = 24,
+    num_epoch: int = 10,
+    batch_size: int = 16,
+    lr: float = 0.0002,
+    betas: Union[float, List[float]] = (0.5, 0.999),
+    device: Union[str, torch.device] = torch.device('cuda'),
+    command: str = 'train',
+    workdir: Union[str,
+                   Path] = Path(__file__).parents[1] / 'workdirs') -> None:
+
+    train_loader, target_loader = init_loader(dataset, data_root, tar_classes,
+                                              batch_size, command)
+    normalizer = norm(dataset, _callable=True)
+
+    surrogate = build_surrogate(surrogate_id, pretrain=True).to(device)
+    surrogate.eval()
+    feat_collecter_handler, feat_collecter = register_collecter(
+        surrogate, midlayer_dict[surrogate_id])
+
+    attack = AIMAttack(device=device)
+    attack.set_mode('train')
+    optim = torch.optim.Adam(attack.get_params(), lr=lr, betas=betas)
+
+    contrastive_loss = ContrastiveLoss(0.2)
+    sim_loss = torch.nn.functional.cosine_similarity
+
+    for epoch in range(1, num_epoch + 1):
+        attack.set_mode('train')
+        enumerator = enumerate(zip(train_loader, target_loader))
+        enumerator = tqdm(enumerator,
+                          total=len(train_loader),
+                          desc=f'Epoch {epoch}')
+        for batch_idx, ((x_nat, y_nat), (x_tar, y_tar)) in enumerator:
+            if torch.any(y_nat == y_tar):
+                continue
+            x_nat, x_tar = x_nat.to(device), x_tar.to(device)
+            y_nat, y_tar = y_nat.to(device), y_tar.to(device)
+            x_adv = attack(x_nat, x_tar)
+
+            logits_nat = surrogate(normalizer(x_nat))
+            feat_nat = feat_collecter.pop()
+            logits_tar = surrogate(normalizer(x_tar))
+            feat_tar = feat_collecter.pop()
+            logits_adv = surrogate(normalizer(x_adv))
+            feat_adv = feat_collecter.pop()
+
+            loss = (contrastive_loss(logits_adv, logits_nat, logits_tar) +
+                    sim_loss(feat_nat, feat_adv) -
+                    sim_loss(feat_tar, feat_adv)).mean()
+
+            optim.zero_grad()
+            loss.backward()
+            optim.step()
+
+    feat_collecter_handler.remove()
+
+    attack.save_ckpt(workdir / 'model.pth')
+
+
+@torch.no_grad()
+def evaluate(
+    ckpt: Union[str, Path],
+    dataset: str,
+    data_root: Union[str, Path],
+    tar_classes: Union[int, List[int]] = 24,
+    batch_size: int = 16,
+    device: Union[str, torch.device] = torch.device('cuda'),
+    command: str = 'train',
+    workdir: Union[str,
+                   Path] = Path(__file__).parents[1] / 'workdirs') -> None:
+    # init dataloader
+    eval_loader, target_loader = init_loader(dataset, data_root, tar_classes,
+                                             batch_size, command)
+    normalizer = norm(dataset, _callable=True)
+    # init attack method
+    attack = AIMAttack(device=device)
+    attack.load_ckpt(ckpt)
+    attack.set_mode('eval')
+    # init evaluate models
+    models = {
+        surrogate_id: build_surrogate(surrogate_id, pretrain=True).to(device)
+        for surrogate_id in list_surrogates()
+    }
+    for surrogate_id in models.keys():
+        models[surrogate_id].eval()
+    model_meters = {
+        surrogate_id: [AverageMeter() for _ in range(2)]
+        for surrogate_id in models.keys()
+    }
+    # evaluate
+    enumerator = enumerate(zip(eval_loader, target_loader))
+    enumerator = tqdm(enumerator, total=len(eval_loader), desc='Eval')
+    for batch_idx, ((x_nat, y_nat), (x_tar, y_tar)) in enumerator:
+        x_nat, y_nat = x_nat.to(device), y_nat.to(device)
+        x_tar, y_tar = x_tar.to(device), y_tar.to(device)
+        x_adv = attack(x_nat, x_tar)
+        for surrogate_id, model in models.items():
+            logits_nat = model(normalizer(x_nat))
+            logits_adv = model(normalizer(x_adv))
+            # collect metrics
+            acc = calc_cls_accuracy(logits_nat, y_nat)
+            asr = calc_cls_accuracy(logits_adv, y_tar)
+            model_meters[surrogate_id][0].update(acc[0].item(), x_nat.size(0))
+            model_meters[surrogate_id][1].update(asr[0].item(), x_nat.size(0))
+    # print result
+    results = {
+        surrogate_id: {
+            'acc': meters[0].avg,
+            'asr': meters[1].avg
+        }
+        for surrogate_id, meters in model_meters.items()
+    }
+    print(pformat(results))
+    with open(workdir / 'results.json', 'w') as f:
+        json.dump(results, f)
+
+
+def main() -> None:
+    args = parse_args()
+    args.command = 'train'
+    args.surrogate_id = 'resnet152'
+    args.num_epoch = 10
+    args.lr = 0.0002
+    args.betas = (0.5, 0.999)
+    if args.command == 'train':
+        train(args.surrogate_id, args.dataset, args.data_root,
+              args.tar_classes, args.num_epoch, args.batch_size, args.lr,
+              args.betas, args.device, args.command, args.workdir)
+    elif args.command == 'evaluate':
+        evaluate(args.ckpt, args.dataset, args.data_root, args.tar_classes,
+                 args.batch_size, args.device, args.command, args.workdir)
+    else:
+        raise NotImplementedError
+
+
+if __name__ == '__main__':
+    main()

attacks/AIM/examples/ens-gen.py

@@ -0,0 +1,523 @@
+#!/usr/bin/env python3
+"""
+Usage
+-----
+./examples/ens-gen.py -d -v -s 0 \
+  --dataset imagenet -b 16 --eps 8 --workdir "workdirs" \
+  --device "cuda:0" \
+  train --n-ep 1 \
+  --surrogate-model-ids vgg19 inception_v3 resnet152 densenet169 \
+  --lr 0.0002 --beta 0.5 0.999 \
+  --use-logit-loss --use-logit-weights --use-logit-softmax-weights
+"""
+import argparse
+import json
+from pathlib import Path
+from pprint import pformat
+from typing import List, Union
+
+import torch
+import torchvision
+from torch.nn import functional as F
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+from gat.datasets import build_dataset, list_datasets
+from gat.datasets.transforms import norm
+from gat.models.attack import CDAAttack
+from gat.models.attack.optim import (SAM, disable_running_stats,
+                                     enable_running_stats)
+from gat.models.surrogate import (build_surrogate, feat_col, list_surrogates,
+                                  midlayer_dict)
+from gat.runtime import AverageMeter, calc_cls_accuracy, fix_random, randid
+
+
+class CLIParser:
+
+    @staticmethod
+    def init_basic_parser(p: argparse.ArgumentParser):
+        g_basic = p.add_argument_group('Basic Settings')
+        g_basic.add_argument('-v',
+                             '--verbose',
+                             action='store_true',
+                             default=False)
+        g_basic.add_argument('-d', '--dev', action='store_true', default=False)
+        g_basic.add_argument('-s', '--seed', type=int, default=0)
+        g_basic.add_argument('--expid', type=str, default=randid(4))
+        g_basic.add_argument('--device', type=str, default='cuda')
+
+        g_path = p.add_argument_group('Path Settings')
+        g_path.add_argument('--workdir', type=str, default='workdirs')
+        g_path.add_argument('--data-root',
+                            type=str,
+                            default=Path(__file__).parent / '../data' /
+                            'in_1k')
+
+        g_ds = p.add_argument_group('Dataset Settings')
+        g_ds.add_argument('--dataset',
+                          type=str,
+                          default='imagenet',
+                          choices=list_datasets())
+        g_ds.add_argument('-b', '--batch-size', type=int, default=16)
+
+        g_at_basic = p.add_argument_group('General Attack Settings')
+        g_at_basic.add_argument('--eps',
+                                '--epsilon',
+                                dest='epsilon',
+                                type=int,
+                                default=8,
+                                choices=[1, 2, 4, 8, 16])
+
+    @staticmethod
+    def post_basic_parser(args: argparse.Namespace):
+        if args.dev:
+            args.workdir = args.workdir.replace('workdirs', 'workdirs-dev')
+        args.workdir = Path(args.workdir) / args.expid
+        args.workdir.mkdir(parents=True, exist_ok=True)
+        args.device = torch.device(args.device)
+
+        args.ckpt = args.workdir / 'model.pth'
+        args.tf_logger = SummaryWriter(args.workdir / 'tf_log')
+
+        args.epsilon /= 255.0
+        if args.command == 'evaluate-pgd':
+            args.alpha /= 255.0
+
+        fix_random(args.seed)
+
+        if args.verbose:
+            print(pformat(vars(args)))
+        with open(args.workdir / f'args-{args.command}.txt', 'w') as f:
+            f.write(pformat(vars(args)))
+
+    @staticmethod
+    def init_train_parser(p: argparse.ArgumentParser):
+        g_at = p.add_argument_group('Attack Settings')
+        g_at.add_argument('--sur-ids',
+                          '--surrogate-model-ids',
+                          dest='surrogate_model_ids',
+                          type=str,
+                          default=['resnet152'],
+                          nargs='+',
+                          choices=list_surrogates())
+        g_at.add_argument('--n-ep',
+                          '--num-epoch',
+                          dest='num_epoch',
+                          type=int,
+                          default=10)
+
+        g_optim = p.add_argument_group('Optimization Settings')
+        g_optim.add_argument('--use-sam', action='store_true', default=False)
+        g_optim.add_argument('--lr', type=float, default=0.0002)
+        g_optim.add_argument('--betas',
+                             type=float,
+                             nargs=2,
+                             default=(0.5, 0.999))
+
+        g_loss = p.add_argument_group('Loss Func Settings')
+        g_loss.add_argument('--use-logit-loss',
+                            action='store_true',
+                            default=False)
+        g_loss.add_argument('--use-logit-kl',
+                            action='store_true',
+                            default=False)
+        g_loss.add_argument('--use-logit-weights',
+                            action='store_true',
+                            default=False)
+        g_loss.add_argument('--use-logit-softmax-weights',
+                            action='store_true',
+                            default=False)
+        g_loss.add_argument('--use-feat-loss',
+                            action='store_true',
+                            default=False)
+        g_loss.add_argument('--use-feat-attn',
+                            action='store_true',
+                            default=False)
+
+    @staticmethod
+    def post_train_parser(args: argparse.Namespace):
+        if args.command == 'train':
+            assert args.use_logit_loss ^ args.use_feat_loss
+            if args.use_logit_kl:
+                assert not args.use_feat_loss
+            if args.use_feat_attn:
+                assert not args.use_logit_loss
+            if args.use_logit_weights:
+                assert args.use_logit_loss
+            if args.use_logit_softmax_weights:
+                assert args.use_logit_loss
+
+    @staticmethod
+    def init_evaluate_parser(p: argparse.ArgumentParser):
+        pass
+
+    @staticmethod
+    def post_evaluate_parser(args: argparse.Namespace):
+        pass
+
+    @staticmethod
+    def init_evaluate_pgd_parser(p: argparse.ArgumentParser):
+        g_at = p.add_argument_group('Attack Settings')
+        g_at.add_argument('--surrogate-model-ids',
+                          type=str,
+                          default=['resnet152'],
+                          nargs='+',
+                          choices=list_surrogates())
+
+        g_optim = p.add_argument_group('Optimization Settings')
+        g_optim.add_argument('--num-step', type=int, default=100)
+        g_optim.add_argument('--alpha',
+                             type=int,
+                             default=2,
+                             choices=[1, 2, 4, 8, 16])
+
+        g_loss = p.add_argument_group('Loss Func Settings')
+        g_loss.add_argument('--use-loss-avg',
+                            action='store_true',
+                            default=False)
+        g_loss.add_argument('--use-logit-avg',
+                            action='store_true',
+                            default=False)
+
+    @staticmethod
+    def post_evaluate_pgd_parser(args: argparse.Namespace):
+        if args.command == 'evaluate-pgd':
+            assert args.use_loss_avg ^ args.use_logit_avg
+
+    @staticmethod
+    def parse_args():
+        p = argparse.ArgumentParser()
+        CLIParser.init_basic_parser(p)
+        sub_p = p.add_subparsers(dest='command')
+
+        CLIParser.init_train_parser(sub_p.add_parser('train'))
+        CLIParser.init_evaluate_parser(sub_p.add_parser('evaluate'))
+        CLIParser.init_evaluate_pgd_parser(sub_p.add_parser('evaluate-pgd'))
+        args = p.parse_args()
+        CLIParser.post_train_parser(args)
+        CLIParser.post_evaluate_parser(args)
+        CLIParser.post_evaluate_pgd_parser(args)
+
+        CLIParser.post_basic_parser(args)
+
+        return args
+
+
+def init_loader(dataset: str,
+                data_root: Union[str, Path],
+                num_epoch: int = 1,
+                batch_size: int = 16,
+                command: str = 'train') -> List[torch.utils.data.DataLoader]:
+    ds = build_dataset(dataset,
+                       data_root=data_root,
+                       is_train=(command == 'train'))
+    dataloader = torch.utils.data.DataLoader(
+        ds,
+        batch_size=batch_size,
+        sampler=torch.utils.data.RandomSampler(ds,
+                                               replacement=True,
+                                               num_samples=len(ds) *
+                                               num_epoch),
+        num_workers=4,
+        pin_memory=True,
+    )
+    normalizer = norm(dataset, _callable=True)
+    return dataloader, normalizer
+
+
+def init_models(model_ids: Union[str, List[str]],
+                device: Union[str, torch.device] = torch.device('cuda')):
+    if isinstance(model_ids, str):
+        model_ids = [model_ids]
+    models = [
+        build_surrogate(_surrogate_id, pretrain=True).to(device)
+        for _surrogate_id in model_ids
+    ]
+    for _ in models:
+        _.eval()
+    return models
+
+
+def calc_loss(x_nat: torch.Tensor,
+              y_nat: torch.Tensor,
+              x_adv: torch.Tensor,
+              feat_collecter: List,
+              surrogate_models: List[torch.nn.Module],
+              normalizer: torchvision.transforms.Compose,
+              use_logit_loss: bool,
+              use_logit_kl: bool,
+              use_logit_weights: bool,
+              use_logit_softmax_weights: bool,
+              use_feat_loss: bool,
+              use_feat_attn: bool,
+              device: Union[str, torch.device] = torch.device('cuda')):
+    loss_sur = []
+    for surrogate_model in surrogate_models:
+        logit_nat = surrogate_model(normalizer(x_nat))
+        feat_nat = feat_collecter.pop()
+        logit_adv = surrogate_model(normalizer(x_adv))
+        feat_adv = feat_collecter.pop()
+        if use_logit_loss:
+            if use_logit_kl:
+                loss_sur.append(-(F.kl_div(F.log_softmax(logit_adv, dim=1),
+                                           F.softmax(logit_nat, dim=1)) +
+                                  F.kl_div(F.log_softmax(logit_nat, dim=1),
+                                           F.softmax(logit_adv, dim=1))))
+            else:
+                loss_sur.append(-(F.cross_entropy(logit_adv, y_nat).mean()))
+        elif use_feat_loss:
+            if use_feat_attn:
+                attn = torch.abs(torch.mean(feat_nat, dim=1, keepdim=True))
+            else:
+                attn = torch.ones_like(feat_nat)
+            loss_sur.append(1 + F.cosine_similarity(attn * feat_nat, attn *
+                                                    feat_adv).mean())
+        else:
+            raise NotImplementedError
+    loss_sur = torch.stack(loss_sur)
+    if use_logit_weights:
+        if use_logit_softmax_weights:
+            loss_weights = torch.nn.functional.softmax(loss_sur)
+        else:
+            loss_weights = torch.nn.functional.softmin(loss_sur)
+        loss_all = torch.sum(loss_weights * loss_sur)
+    else:
+        loss_all = loss_sur.mean()
+
+    return loss_all
+
+
+def train(surrogate_model_ids: Union[str, List[str]],
+          epsilon: float = 16.0 / 255.0,
+          num_epoch: int = 10,
+          dataset: str = 'imagenet',
+          batch_size: int = 16,
+          use_sam: bool = False,
+          lr: float = 0.0002,
+          betas: Union[float, List[float]] = (0.5, 0.999),
+          use_logit_loss: bool = False,
+          use_logit_kl: bool = False,
+          use_logit_weights: bool = False,
+          use_logit_softmax_weights: bool = False,
+          use_feat_loss: bool = False,
+          use_feat_attn: bool = False,
+          device: Union[str, torch.device] = torch.device('cuda'),
+          workdir: Union[str, Path] = Path(__file__).parents[1] / 'workdirs',
+          data_root: Union[str,
+                           Path] = Path(__file__).parent / '../data' / 'in_1k',
+          tf_logger: SummaryWriter = None) -> None:
+    """
+    Train the attack model with the given surrogate models.
+    """
+    loader, normalizer = init_loader(dataset, data_root, num_epoch, batch_size,
+                                     'train')
+    surrogate_models = init_models(surrogate_model_ids, device)
+
+    attack = CDAAttack(device=device, epsilon=epsilon)
+    attack.set_mode('train')
+    if use_sam:
+        optim = SAM(attack.get_params(), torch.optim.Adam, lr=lr, betas=betas)
+    else:
+        optim = torch.optim.Adam(attack.get_params(), lr=lr, betas=betas)
+
+    with feat_col(surrogate_models,
+                  [midlayer_dict[_]
+                   for _ in surrogate_model_ids]) as feat_collecter:
+        attack.set_mode('train')
+        enumerator = tqdm(enumerate(loader), total=len(loader), desc='')
+        for step, (x_nat, y_nat) in enumerator:
+            x_nat, y_nat = x_nat.to(device), y_nat.to(device)
+
+            if use_sam:
+                # 1
+                enable_running_stats(attack.get_model())
+                loss_v = calc_loss(x_nat, y_nat, attack(x_nat), feat_collecter,
+                                   surrogate_models, normalizer,
+                                   use_logit_loss, use_logit_kl,
+                                   use_logit_weights,
+                                   use_logit_softmax_weights, use_feat_loss,
+                                   use_feat_attn, device)
+                loss_v.backward()
+                optim.first_step(zero_grad=True)
+                # 2
+                disable_running_stats(attack.get_model())
+                calc_loss(x_nat, y_nat, attack(x_nat), feat_collecter,
+                          surrogate_models, normalizer, use_logit_loss,
+                          use_logit_kl, use_logit_weights,
+                          use_logit_softmax_weights, use_feat_loss,
+                          use_feat_attn, device).backward()
+                optim.second_step(zero_grad=True)
+            else:
+                x_adv = attack(x_nat)
+                loss_v = calc_loss(x_nat, y_nat, x_adv, feat_collecter,
+                                   surrogate_models, normalizer,
+                                   use_logit_loss, use_logit_kl,
+                                   use_logit_weights,
+                                   use_logit_softmax_weights, use_feat_loss,
+                                   use_feat_attn, device)
+                optim.zero_grad()
+                loss_v.backward()
+                optim.step()
+
+            if tf_logger:
+                tf_logger.add_scalar('loss', loss_v.item(), step)
+                tf_logger.add_scalar('lr', optim.param_groups[0]['lr'], step)
+
+    attack.save_ckpt(workdir / 'model.pth')
+
+
+@torch.no_grad()
+def evaluate(
+    ckpt: Union[str, Path],
+    epsilon: float = 16.0 / 255.0,
+    dataset: str = 'imagenet',
+    batch_size: int = 16,
+    device: Union[str, torch.device] = torch.device('cuda'),
+    workdir: Union[str, Path] = Path(__file__).parents[1] / 'workdirs',
+    data_root: Union[str, Path] = Path(__file__).parent / '../data' / 'in_1k',
+) -> None:
+    """
+    Evaluate the attack model with the given surrogate models
+    """
+    loader, normalizer = init_loader(dataset, data_root, 1, batch_size,
+                                     'evaluate')
+    target_models = {
+        k: v
+        for k, v in zip(list_surrogates(),
+                        init_models(list_surrogates(), device))
+    }
+    target_acc_meters = {
+        target_model_id: [AverageMeter() for _ in range(2)]
+        for target_model_id in target_models.keys()
+    }
+    # init attack method
+    attack = CDAAttack(device=device, epsilon=epsilon)
+    attack.load_ckpt(ckpt)
+    attack.set_mode('eval')
+    # evaluate
+    enumerator = tqdm(enumerate(loader), total=len(loader), desc='Eval')
+    for step, (x_nat, y_nat) in enumerator:
+        x_nat, y_nat = x_nat.to(device), y_nat.to(device)
+        x_adv = attack(x_nat)
+        for target_model_id, target_model in target_models.items():
+            logit_nat = target_model(normalizer(x_nat))
+            logit_adv = target_model(normalizer(x_adv))
+            # collect metrics
+            target_acc = calc_cls_accuracy(logit_nat, y_nat)
+            target_asr = calc_cls_accuracy(logit_adv, y_nat)
+            target_acc_meters[target_model_id][0].update(
+                target_acc[0].item(), x_nat.size(0))
+            target_acc_meters[target_model_id][1].update(
+                target_asr[0].item(), x_nat.size(0))
+    results = {
+        target_model_id: {
+            'nat_acc': target_acc_meter[0].avg,
+            'adv_acc': target_acc_meter[1].avg
+        }
+        for target_model_id, target_acc_meter in target_acc_meters.items()
+    }
+    print(pformat(results))
+    with open(workdir / 'results.json', 'w') as f:
+        json.dump(results, f)
+
+
+def evaluate_pgd(
+    surrogate_model_ids: Union[str, List[str]],
+    epsilon: float = 16.0 / 255.0,
+    num_step: int = 1000,
+    alpha: float = 2.0 / 255.0,
+    dataset: str = 'imagenet',
+    batch_size: int = 16,
+    use_loss_avg: bool = False,
+    use_logit_avg: bool = False,
+    device: Union[str, torch.device] = torch.device('cuda'),
+    workdir: Union[str, Path] = Path(__file__).parents[1] / 'workdirs',
+    data_root: Union[str, Path] = Path(__file__).parent / '../data' / 'in_1k',
+):
+    loader, normalizer = init_loader(dataset, data_root, 1, batch_size,
+                                     'evaluate')
+    surrogate_models = init_models(surrogate_model_ids, device)
+    target_models = {
+        k: v
+        for k, v in zip(list_surrogates(),
+                        init_models(list_surrogates(), device))
+    }
+    target_acc_meters = {
+        target_model_id: [AverageMeter() for _ in range(2)]
+        for target_model_id in target_models.keys()
+    }
+    # evaluate
+    enumerator = tqdm(enumerate(loader), total=len(loader), desc='')
+    for step, (x_nat, y_nat) in enumerator:
+        x_nat, y_nat = x_nat.to(device), y_nat.to(device)
+        # attack
+        x_nat_ori = x_nat.data
+        for _ in range(num_step):
+            x_nat.requires_grad = True
+            if use_loss_avg:
+                loss_all = 0.0
+                for surrogate_model in surrogate_models:
+                    logit = surrogate_model(x_nat)
+                    surrogate_model.zero_grad()
+                    loss_all += F.cross_entropy(logit, y_nat)
+            elif use_logit_avg:
+                logit = torch.stack([
+                    surrogate_model(x_nat)
+                    for surrogate_model in surrogate_models
+                ]).mean(dim=0)
+                loss_all = F.cross_entropy(logit, y_nat)
+            else:
+                raise NotADirectoryError
+            loss_all.backward()
+            x_adv_ = x_nat + alpha * x_nat.grad.sign()
+            eta = torch.clamp(x_adv_ - x_nat_ori, min=-epsilon, max=epsilon)
+            x_nat = torch.clamp(x_nat_ori + eta, min=0.0, max=1.0).detach_()
+        x_adv = x_nat
+        x_nat = x_nat_ori
+        # eval
+        with torch.no_grad():
+            for target_model_id, target_model in target_models.items():
+                logit_nat = target_model(normalizer(x_nat))
+                logit_adv = target_model(normalizer(x_adv))
+                # collect
+                target_acc_ = calc_cls_accuracy(logit_nat, y_nat)
+                target_asr_ = calc_cls_accuracy(logit_adv, y_nat)
+                target_acc_meters[target_model_id][0].update(
+                    target_acc_[0].item(), x_nat.size(0))
+                target_acc_meters[target_model_id][1].update(
+                    target_asr_[0].item(), x_nat.size(0))
+    results = {
+        target_model_id: {
+            'nat_acc': target_acc_meter[0].avg,
+            'adv_acc': target_acc_meter[1].avg
+        }
+        for target_model_id, target_acc_meter in target_acc_meters.items()
+    }
+    print(pformat(results))
+    with open(workdir / 'results-pgd.json', 'w') as f:
+        json.dump(results, f)
+
+
+def main() -> None:
+    args = CLIParser.parse_args()
+    if args.command == 'train':
+        train(args.surrogate_model_ids, args.epsilon, args.num_epoch,
+              args.dataset, args.batch_size, args.use_sam, args.lr, args.betas,
+              args.use_logit_loss, args.use_logit_kl, args.use_logit_weights,
+              args.use_logit_softmax_weights, args.use_feat_loss,
+              args.use_feat_attn, args.device, args.workdir, args.data_root,
+              args.tf_logger)
+    elif args.command == 'evaluate':
+        evaluate(args.ckpt, args.epsilon, args.dataset, args.batch_size,
+                 args.device, args.workdir, args.data_root)
+    elif args.command == 'evaluate-pgd':
+        evaluate_pgd(args.surrogate_model_ids, args.epsilon, args.num_step,
+                     args.alpha, args.dataset, args.batch_size,
+                     args.use_loss_avg, args.use_logit_avg, args.device,
+                     args.workdir, args.data_root)
+    else:
+        raise NotImplementedError
+
+
+if __name__ == '__main__':
+    main()

attacks/AIM/examples/workdirs/1MsK/args.txt

@@ -0,0 +1,11 @@
+{'batch_size': 16,
+ 'ckpt': WindowsPath('workdirs/1MsK/model.pth'),
+ 'command': None,
+ 'data_root': WindowsPath('D:/Sharing/Programs/3-Durable-Adv/PyCIL-master/attacks/GAT/examples/../data/in_1k'),
+ 'dataset': 'imagenet',
+ 'device': device(type='cuda'),
+ 'expid': '1MsK',
+ 'seed': 0,
+ 'tar_classes': 24,
+ 'verbose': False,
+ 'workdir': WindowsPath('workdirs/1MsK')}

attacks/AIM/examples/workdirs/CGKQ/args.txt

@@ -0,0 +1,11 @@
+{'batch_size': 16,
+ 'ckpt': WindowsPath('workdirs/CGKQ/model.pth'),
+ 'command': None,
+ 'data_root': WindowsPath('D:/Sharing/Programs/3-Durable-Adv/PyCIL-master/attacks/GAT/examples/../data/in_1k'),
+ 'dataset': 'imagenet',
+ 'device': device(type='cuda'),
+ 'expid': 'CGKQ',
+ 'seed': 0,
+ 'tar_classes': 24,
+ 'verbose': False,
+ 'workdir': WindowsPath('workdirs/CGKQ')}

attacks/AIM/examples/workdirs/Fvu1/args.txt

@@ -0,0 +1,11 @@
+{'batch_size': 16,
+ 'ckpt': WindowsPath('workdirs/Fvu1/model.pth'),
+ 'command': None,
+ 'data_root': WindowsPath('D:/Sharing/Programs/3-Durable-Adv/PyCIL-master/attacks/GAT/examples/../data/in_1k'),
+ 'dataset': 'imagenet',
+ 'device': device(type='cuda'),
+ 'expid': 'Fvu1',
+ 'seed': 0,
+ 'tar_classes': 24,
+ 'verbose': False,
+ 'workdir': WindowsPath('workdirs/Fvu1')}

attacks/AIM/examples/workdirs/Qonx/args.txt

@@ -0,0 +1,11 @@
+{'batch_size': 16,
+ 'ckpt': WindowsPath('workdirs/Qonx/model.pth'),
+ 'command': None,
+ 'data_root': WindowsPath('D:/Sharing/Programs/3-Durable-Adv/PyCIL-master/attacks/GAT/examples/../data/in_1k'),
+ 'dataset': 'imagenet',
+ 'device': device(type='cuda'),
+ 'expid': 'Qonx',
+ 'seed': 0,
+ 'tar_classes': 24,
+ 'verbose': False,
+ 'workdir': WindowsPath('workdirs/Qonx')}

attacks/AIM/examples/workdirs/fnNs/args.txt

@@ -0,0 +1,11 @@
+{'batch_size': 16,
+ 'ckpt': WindowsPath('workdirs/fnNs/model.pth'),
+ 'command': None,
+ 'data_root': WindowsPath('D:/Sharing/Programs/3-Durable-Adv/PyCIL-master/attacks/GAT/examples/../data/in_1k'),
+ 'dataset': 'imagenet',
+ 'device': device(type='cuda'),
+ 'expid': 'fnNs',
+ 'seed': 0,
+ 'tar_classes': 24,
+ 'verbose': False,
+ 'workdir': WindowsPath('workdirs/fnNs')}

attacks/AIM/examples/workdirs/jtMb/args.txt

@@ -0,0 +1,11 @@
+{'batch_size': 16,
+ 'ckpt': WindowsPath('workdirs/jtMb/model.pth'),
+ 'command': None,
+ 'data_root': WindowsPath('D:/Sharing/Programs/3-Durable-Adv/PyCIL-master/attacks/GAT/examples/../data/in_1k'),
+ 'dataset': 'imagenet',
+ 'device': device(type='cuda'),
+ 'expid': 'jtMb',
+ 'seed': 0,
+ 'tar_classes': 24,
+ 'verbose': False,
+ 'workdir': WindowsPath('workdirs/jtMb')}

attacks/AIM/examples/workdirs/krhX/args.txt

@@ -0,0 +1,11 @@
+{'batch_size': 16,
+ 'ckpt': WindowsPath('workdirs/krhX/model.pth'),
+ 'command': None,
+ 'data_root': WindowsPath('D:/Sharing/Programs/3-Durable-Adv/PyCIL-master/attacks/GAT/examples/../data/in_1k'),
+ 'dataset': 'imagenet',
+ 'device': device(type='cuda'),
+ 'expid': 'krhX',
+ 'seed': 0,
+ 'tar_classes': 24,
+ 'verbose': False,
+ 'workdir': WindowsPath('workdirs/krhX')}

attacks/AIM/setup.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+import sys
+from pathlib import Path
+from typing import Union
+
+from setuptools import find_packages, setup
+
+try:
+    sys.path.append(str(Path(__file__).parent / 'src'))
+    import gat
+except ImportError:
+    print('Please install aim_attack package')
+    sys.exit(1)
+
+
+def read(path: Union[str, Path]) -> str:
+    with open(path, 'r') as f:
+        return f.read()
+
+
+setup(name='pygat',
+      version=gat.VERSION,
+      description='GAT: Generative Attack Toolbox',
+      long_description=read('README.md'),
+      long_description_content_type='text/markdown',
+      author='Terry Li',
+      url='https://terrytengli.com/GAT/',
+      classifiers=[
+          'Programming Language :: Python :: 3',
+          'License :: OSI Approved :: MIT License'
+      ],
+      python_requires='>=3.10',
+      packages=find_packages('src'),
+      package_dir={'': 'src'},
+      entry_points={
+          'console_scripts': [
+              'aim-api=gat.runtime.api.aim_attack:main',
+          ],
+      },
+      install_requires=[
+          'tqdm', 'tabulate', 'torch', 'torchvision', 'tensorboard', 'jupyter'
+      ],
+      extras_require={
+          'cli': ['python-multipart', 'fastapi', 'uvicorn'],
+          'ens-gen': []
+      },
+      include_package_data=True)

attacks/AIM/src/gat/__init__.py

@@ -0,0 +1 @@
+VERSION = '202411.2'

attacks/AIM/src/gat/datasets/__init__.py

@@ -0,0 +1,5 @@
+from .builder import build_dataset, list_datasets
+from .cub import cub
+from .imagenet import imagenet
+
+__all__ = ['build_dataset', 'list_datasets', 'imagenet', 'cub']

attacks/AIM/src/gat/datasets/builder.py

@@ -0,0 +1,11 @@
+from ..runtime.factory import Registry
+
+DATASET_REGISTRY = Registry('DATASET')
+
+
+def build_dataset(_type: str, *args, **kwargs) -> object:
+    return DATASET_REGISTRY.get(_type)(*args, **kwargs)
+
+
+def list_datasets():
+    return list(DATASET_REGISTRY._obj_map.keys())

attacks/AIM/src/gat/datasets/cub.py

@@ -0,0 +1,6 @@
+from .builder import DATASET_REGISTRY
+
+
+@DATASET_REGISTRY.register()
+def cub():
+    raise NotADirectoryError

attacks/AIM/src/gat/datasets/env.py

@@ -0,0 +1,6 @@
+CIFAR10_DEFAULT_MEAN = (0.4914, 0.4822, 0.4465)
+CIFAR10_DEFAULT_STD = (0.2470, 0.2435, 0.2616)
+CIFAR100_DEFAULT_MEAN = (0.5071, 0.4865, 0.4409)
+CIFAR100_DEFAULT_STD = (0.2673, 0.2564, 0.2762)
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)

attacks/AIM/src/gat/datasets/imagenet.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+from typing import List, Union
+
+import torch
+import torchvision
+
+from .builder import DATASET_REGISTRY
+from .transforms import resize_256_224, to_color, to_ts
+
+
+@DATASET_REGISTRY.register()
+def imagenet(
+    data_root: Union[str, Path],
+    is_train: bool = True,
+    filter_class: Union[int, List[int]] = None,
+) -> torch.utils.data.Dataset:
+    if isinstance(data_root, str):
+        data_root = Path(data_root)
+    if is_train:
+        data_root = data_root / 'train'
+    else:
+        data_root = data_root / 'val'
+
+    _transforms = resize_256_224() + to_color() + to_ts()
+
+    _ds = torchvision.datasets.ImageFolder(
+        data_root,
+        transform=torchvision.transforms.Compose(_transforms),
+    )
+
+    if isinstance(filter_class, int):
+        filter_class = [filter_class]
+    if filter_class:
+        _ds.samples = list(filter(lambda x: x[1] in filter_class, _ds.samples))
+
+    return _ds

attacks/AIM/src/gat/datasets/transforms.py

@@ -0,0 +1,55 @@
+from typing import List
+
+import torchvision
+
+from . import env
+
+
+def resize_256_224() -> List:
+    return [
+        torchvision.transforms.Resize(size=256),
+        torchvision.transforms.CenterCrop(size=(224, 224)),
+    ]
+
+
+def resize_512_448() -> List:
+    return [
+        torchvision.transforms.Resize(size=512),
+        torchvision.transforms.CenterCrop(size=(448, 448)),
+    ]
+
+
+def resize_224() -> List:
+    return [torchvision.transforms.Resize(size=224)]
+
+
+def hflip(p: float = 0.5) -> List:
+    assert 0 <= p <= 1
+    return [torchvision.transforms.RandomHorizontalFlip(p)]
+
+
+def to_ts() -> List:
+    return [torchvision.transforms.ToTensor()]
+
+
+def to_pil() -> List:
+    return [torchvision.transforms.ToPILImage()]
+
+
+def to_color() -> List:
+    return [
+        torchvision.transforms.Lambda(lambda x: x.convert('RGB')
+                                      if x.mode != 'RGB' else x)
+    ]
+
+
+def norm(dataset: str = 'IMAGENET', _callable: bool = False) -> List:
+    dataset = dataset.upper()
+    mean_std = (
+        getattr(env, dataset + '_DEFAULT_MEAN'),
+        getattr(env, dataset + '_DEFAULT_STD'),
+    )
+    transforms = [torchvision.transforms.Normalize(*mean_std)]
+    if _callable:
+        return torchvision.transforms.Compose(transforms)
+    return transforms

attacks/AIM/src/gat/models/attack/__init__.py

@@ -0,0 +1,5 @@
+from .aim_attack import AIMAttack
+from .cda_attack import CDAAttack
+from .loss.logits import ContrastiveLoss
+
+__all__ = ['AIMAttack', 'CDAAttack', 'ContrastiveLoss']

attacks/AIM/src/gat/models/attack/aim_attack.py

@@ -0,0 +1,14 @@
+import torch
+
+from .base_attack import BaseGenerativeAttack
+from .generator.aim import AIMGenerator
+
+
+class AIMAttack(BaseGenerativeAttack):
+
+    def set_adv_gen(self):
+        self.adv_gen = AIMGenerator().to(self.device)
+
+    def attack(self, x_nat, *extra_inputs) -> torch.Tensor:
+        x_guid = extra_inputs[0].to(self.device)
+        return self.adv_gen(x_nat, x_guid)

attacks/AIM/src/gat/models/attack/base_attack.py

@@ -0,0 +1,60 @@
+import abc
+from collections import OrderedDict
+from pathlib import Path
+from typing import Union
+
+import torch
+
+
+class BaseGenerativeAttack(abc.ABC):
+
+    def __init__(self,
+                 device: Union[str, torch.device],
+                 epsilon: float = 32 / 255) -> None:
+        if isinstance(device, str):
+            device = torch.device(device)
+        self.device = device
+        self.set_adv_gen()
+        self.set_mode('eval')
+        self.epsilon = epsilon
+
+    @abc.abstractmethod
+    def set_adv_gen(self):
+        pass
+
+    def load_ckpt(self, ckpt: Union[str, Path, OrderedDict]) -> None:
+        if isinstance(ckpt, str):
+            ckpt = Path(ckpt)
+        if isinstance(ckpt, Path):
+            if not ckpt.exists():
+                raise FileNotFoundError(f'File not found: {ckpt}')
+            ckpt = torch.load(ckpt, map_location=self.device)
+        self.adv_gen.load_state_dict(ckpt)
+        self.adv_gen.to(self.device)
+
+    def save_ckpt(self, ckpt: Union[str, Path]) -> None:
+        if isinstance(ckpt, str):
+            ckpt = Path(ckpt)
+        _adv_gen_cpu = self.adv_gen.to('cpu')
+        torch.save(_adv_gen_cpu.state_dict(), ckpt)
+
+    def get_params(self) -> torch.nn.Parameter:
+        return self.adv_gen.parameters()
+
+    def get_model(self) -> torch.nn.Module:
+        return self.adv_gen
+
+    def set_mode(self, mode: str) -> None:
+        assert mode in ['train', 'eval']
+        self.adv_gen.train() if mode == 'train' else self.adv_gen.eval()
+
+    @abc.abstractmethod
+    def attack(self, *args) -> torch.Tensor:
+        pass
+
+    def __call__(self, x_nat: torch.Tensor, *extra_inputs) -> torch.Tensor:
+        x_adv = self.attack(x_nat, *extra_inputs)
+        x_adv = torch.min(torch.max(x_adv, x_nat - self.epsilon),
+                          x_nat + self.epsilon)
+        torch.clamp_(x_adv, 0.0, 1.0)
+        return x_adv

attacks/AIM/src/gat/models/attack/cda_attack.py

@@ -0,0 +1,13 @@
+import torch
+
+from .base_attack import BaseGenerativeAttack
+from .generator.cda import CDAGenerator
+
+
+class CDAAttack(BaseGenerativeAttack):
+
+    def set_adv_gen(self):
+        self.adv_gen = CDAGenerator().to(self.device)
+
+    def attack(self, x_nat, *extra_inputs) -> torch.Tensor:
+        return self.adv_gen(x_nat)

attacks/AIM/src/gat/models/attack/generator/aim.py

@@ -0,0 +1,179 @@
+import torch
+from torch import nn
+
+
+class EnhancedBN(nn.Module):
+    def __init__(self, nc: int, sty_nc: int = 3, sty_nhidden: int = 128):
+        super(EnhancedBN, self).__init__()
+        self.bn = nn.BatchNorm2d(nc)
+        self.mapping = nn.Conv2d(
+            in_channels=sty_nc,
+            out_channels=sty_nhidden,
+            kernel_size=3,
+            padding=1,
+            stride=1,
+        )
+        self.gamma = nn.Conv2d(
+            in_channels=sty_nhidden,
+            out_channels=nc,
+            kernel_size=3,
+            padding=1,
+            stride=1,
+        )
+        self.beta = nn.Conv2d(
+            in_channels=sty_nhidden,
+            out_channels=nc,
+            kernel_size=3,
+            padding=1,
+            stride=1,
+        )
+        self.init_weight()
+
+    def init_weight(self):
+        nn.init.kaiming_normal_(self.mapping.weight)
+        nn.init.kaiming_normal_(self.gamma.weight)
+        nn.init.kaiming_normal_(self.beta.weight)
+
+    def forward(self, base, sty):
+        bn = self.bn(base)
+        sty_resized = torch.nn.functional.interpolate(
+            sty, size=bn.size()[2:], mode='bilinear'
+        )
+        actv = torch.nn.functional.relu(self.mapping(sty_resized))
+        # style injection
+        bn = bn * (1 + self.gamma(actv)) + self.beta(actv)
+        return bn
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, num_filters):
+        super(ResidualBlock, self).__init__()
+        self.block1 = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(
+                in_channels=num_filters,
+                out_channels=num_filters,
+                kernel_size=3,
+                stride=1,
+                padding=0,
+                bias=False,
+            ),
+        )
+        self.bn1 = EnhancedBN(num_filters)
+        self.block2 = nn.Sequential(
+            nn.ReLU(True),
+            nn.Dropout(0.5),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(
+                in_channels=num_filters,
+                out_channels=num_filters,
+                kernel_size=3,
+                stride=1,
+                padding=0,
+                bias=False,
+            ),
+        )
+        self.bn2 = EnhancedBN(num_filters)
+
+    def forward(self, x, sty):
+        residual = self.block1(x)
+        residual = self.bn1(residual, sty)
+        residual = self.block2(residual)
+        residual = self.bn2(residual, sty)
+        return x + residual
+
+
+ngf = 64
+
+
+class ResNetGenerator(nn.Module):
+    def __init__(self):
+        super(ResNetGenerator, self).__init__()
+        self.block1 = nn.Sequential(
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(3, ngf, kernel_size=7, padding=0, bias=False),
+        )
+        self.bn1 = EnhancedBN(ngf)
+        # Input size = 3, n, n
+        self.block2 = nn.Sequential(
+            nn.Conv2d(
+                ngf, ngf * 2, kernel_size=3, stride=2, padding=1, bias=False
+            ),
+        )
+        self.bn2 = EnhancedBN(ngf * 2)
+        # Input size = 3, n/2, n/2
+        self.block3 = nn.Sequential(
+            nn.Conv2d(
+                ngf * 2,
+                ngf * 4,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                bias=False,
+            ),
+        )
+        self.bn3 = EnhancedBN(ngf * 4)
+        # Input size = 3, n/4, n/4
+        # Residual Blocks: 6
+        self.resblock1 = ResidualBlock(ngf * 4)
+        self.resblock2 = ResidualBlock(ngf * 4)
+        self.resblock3 = ResidualBlock(ngf * 4)
+        self.resblock4 = ResidualBlock(ngf * 4)
+        self.resblock5 = ResidualBlock(ngf * 4)
+        self.resblock6 = ResidualBlock(ngf * 4)
+        # Input size = 3, n/4, n/4
+        self.upsampl1 = nn.ConvTranspose2d(
+            ngf * 4,
+            ngf * 2,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            output_padding=1,
+            bias=False,
+        )
+        self.ubn1 = EnhancedBN(ngf * 2)
+        # Input size = 3, n/2, n/2
+        self.upsampl2 = nn.ConvTranspose2d(
+            ngf * 2,
+            ngf,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+            output_padding=1,
+            bias=False,
+        )
+        self.ubn2 = EnhancedBN(ngf)
+        # Input size = 3, n, n
+        self.blockf = nn.Sequential(
+            nn.ReflectionPad2d(3), nn.Conv2d(ngf, 3, kernel_size=7, padding=0)
+        )
+
+    def forward(self, input, sty):
+        x = self.block1(input)
+        x = self.bn1(x, sty)
+        x = torch.nn.functional.relu(x)
+        x = self.block2(x)
+        x = self.bn2(x, sty)
+        x = torch.nn.functional.relu(x)
+        x = self.block3(x)
+        x = self.bn3(x, sty)
+        x = torch.nn.functional.relu(x)
+        # =============================
+        x = self.resblock1(x, sty)
+        x = self.resblock2(x, sty)
+        x = self.resblock3(x, sty)
+        x = self.resblock4(x, sty)
+        x = self.resblock5(x, sty)
+        x = self.resblock6(x, sty)
+        # =============================
+        x = self.upsampl1(x)
+        x = self.ubn1(x, sty)
+        x = torch.nn.functional.relu(x)
+        x = self.upsampl2(x)
+        x = self.ubn2(x, sty)
+        x = torch.nn.functional.relu(x)
+        x = self.blockf(x)
+        return (torch.tanh(x) + 1) / 2
+
+
+AIMGenerator = ResNetGenerator

attacks/AIM/src/gat/models/attack/generator/cda.py

@@ -0,0 +1,146 @@
+# Code is copied from
+# github.com/Alibaba-AAIG/Beyond-ImageNet-Attack:generator.py@863b7
+import torch
+from torch import nn
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, num_filters):
+        super(ResidualBlock, self).__init__()
+        self.block = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(
+                in_channels=num_filters,
+                out_channels=num_filters,
+                kernel_size=3,
+                stride=1,
+                padding=0,
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_filters),
+            nn.ReLU(True),
+            nn.Dropout(0.5),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(
+                in_channels=num_filters,
+                out_channels=num_filters,
+                kernel_size=3,
+                stride=1,
+                padding=0,
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_filters),
+        )
+
+    def forward(self, x):
+        residual = self.block(x)
+        return x + residual
+
+
+ngf = 64
+
+
+class ResNetGenerator(nn.Module):
+    """
+    https://github.com/Alibaba-AAIG/Beyond-ImageNet-Attack/blob/863b758ee4f4a6d3d4e7777c5f94f457fa449f73/generator.py#L14
+
+    Test Case:
+    >>> netG = ResNetGenerator()
+    >>> test_sample = torch.rand(1, 3, 224, 224)
+    >>> print("Generator output:", netG(test_sample).size())
+    >>> print(
+    >>>     "Generator parameters:",
+    >>>     sum(p.numel() for p in netG.parameters() if p.requires_grad),
+    >>> )
+    """
+
+    def __init__(self, inception=False):
+        super(ResNetGenerator, self).__init__()
+        self.inception = inception
+        self.block1 = nn.Sequential(
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(3, ngf, kernel_size=7, padding=0, bias=False),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True),
+        )
+        # output: (ngf) x (n) x (n)
+        self.block2 = nn.Sequential(
+            nn.Conv2d(
+                ngf, ngf * 2, kernel_size=3, stride=2, padding=1, bias=False
+            ),
+            nn.BatchNorm2d(ngf * 2),
+            nn.ReLU(True),
+        )
+        # output: (ngf*2) x (n/2) x (n/2)
+        self.block3 = nn.Sequential(
+            nn.Conv2d(
+                ngf * 2,
+                ngf * 4,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(ngf * 4),
+            nn.ReLU(True),
+        )
+        # output: (ngf*4) x (n/4) x (n/4)
+        self.resblock1 = ResidualBlock(ngf * 4)
+        self.resblock2 = ResidualBlock(ngf * 4)
+        self.resblock3 = ResidualBlock(ngf * 4)
+        self.resblock4 = ResidualBlock(ngf * 4)
+        self.resblock5 = ResidualBlock(ngf * 4)
+        self.resblock6 = ResidualBlock(ngf * 4)
+        # output: (ngf*4) x (n/4) x (n/4)
+        self.upsampl1 = nn.Sequential(
+            nn.ConvTranspose2d(
+                ngf * 4,
+                ngf * 2,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                output_padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(ngf * 2),
+            nn.ReLU(True),
+        )
+        # output: (ngf*2) x (n/2) x (n/2)
+        self.upsampl2 = nn.Sequential(
+            nn.ConvTranspose2d(
+                ngf * 2,
+                ngf,
+                kernel_size=3,
+                stride=2,
+                padding=1,
+                output_padding=1,
+                bias=False,
+            ),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True),
+        )
+        # output: (ngf) x (n) x (n)
+        self.blockf = nn.Sequential(
+            nn.ReflectionPad2d(3), nn.Conv2d(ngf, 3, kernel_size=7, padding=0)
+        )
+        self.crop = nn.ConstantPad2d((0, -1, -1, 0), 0)
+
+    def forward(self, input):
+        x = self.block1(input)
+        x = self.block2(x)
+        x = self.block3(x)
+        x = self.resblock1(x)
+        x = self.resblock2(x)
+        x = self.resblock3(x)
+        x = self.resblock4(x)
+        x = self.resblock5(x)
+        x = self.resblock6(x)
+        x = self.upsampl1(x)
+        x = self.upsampl2(x)
+        x = self.blockf(x)
+        if self.inception:
+            x = self.crop(x)
+        return (torch.tanh(x) + 1) / 2
+
+
+CDAGenerator = ResNetGenerator

attacks/AIM/src/gat/models/attack/loss/logits.py

@@ -0,0 +1,27 @@
+import torch
+
+
+class ContrastiveLoss(torch.nn.Module):
+    """
+    Contrastive loss
+    Adapted from: (OnlineContrastiveLoss)
+        https://github.com/adambielski/siamese-triplet/blob/master/losses.py
+    Based on: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf
+    """
+
+    def __init__(self, margin):
+        super(ContrastiveLoss, self).__init__()
+        self.margin = margin
+
+    def forward(self, anchors, negatives, positives):
+        anchors = anchors / anchors.norm(dim=-1, keepdim=True)
+        negatives = negatives / negatives.norm(dim=-1, keepdim=True)
+        positives = positives / positives.norm(dim=-1, keepdim=True)
+
+        positive_loss = (anchors - positives).pow(2).sum(1)
+        negative_loss = torch.nn.functional.relu(
+            self.margin - (anchors - negatives).pow(2).sum(1).sqrt()).pow(2)
+
+        loss = 0.5 * torch.cat([positive_loss, negative_loss], dim=0)
+
+        return loss.mean()

attacks/AIM/src/gat/models/attack/optim/__init__.py

@@ -0,0 +1,3 @@
+from .sam import SAM, disable_running_stats, enable_running_stats
+
+__all__ = ['SAM', 'enable_running_stats', 'disable_running_stats']

attacks/AIM/src/gat/models/attack/optim/sam.py

@@ -0,0 +1,100 @@
+# Code is copied from
+# github.com/davda54/sam/sam.py@3c3afdb
+import torch
+from torch.nn.modules.batchnorm import _BatchNorm
+
+
+class SAM(torch.optim.Optimizer):
+
+    def __init__(self,
+                 params,
+                 base_optimizer,
+                 rho=0.05,
+                 adaptive=False,
+                 **kwargs):
+        assert rho >= 0.0, f"Invalid rho, should be non-negative: {rho}"
+
+        defaults = dict(rho=rho, adaptive=adaptive, **kwargs)
+        super(SAM, self).__init__(params, defaults)
+
+        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
+        self.param_groups = self.base_optimizer.param_groups
+        self.defaults.update(self.base_optimizer.defaults)
+
+    @torch.no_grad()
+    def first_step(self, zero_grad=False):
+        grad_norm = self._grad_norm()
+        for group in self.param_groups:
+            scale = group['rho'] / (grad_norm + 1e-12)
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                self.state[p]['old_p'] = p.data.clone()
+                e_w = (torch.pow(p, 2)
+                       if group['adaptive'] else 1.0) * p.grad * scale.to(p)
+                p.add_(e_w)  # climb to the local maximum "w + e(w)"
+
+        if zero_grad:
+            self.zero_grad()
+
+    @torch.no_grad()
+    def second_step(self, zero_grad=False):
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                p.data = self.state[p][
+                    'old_p']  # get back to "w" from "w + e(w)"
+
+        self.base_optimizer.step()  # do the actual "sharpness-aware" update
+
+        if zero_grad:
+            self.zero_grad()
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        assert closure is not None, \
+                'Sharpness Aware Minimization requires closure, ' \
+                'but it was not provided'
+        closure = torch.enable_grad()(
+            closure)  # the closure should do a full forward-backward pass
+
+        self.first_step(zero_grad=True)
+        closure()
+        self.second_step()
+
+    def _grad_norm(self):
+        # put everything on the same device, in case of model parallelism
+        shared_device = self.param_groups[0]['params'][0].device
+        norm = torch.norm(torch.stack([
+            ((torch.abs(p) if group['adaptive'] else 1.0) *
+             p.grad).norm(p=2).to(shared_device) for group in self.param_groups
+            for p in group['params'] if p.grad is not None
+        ]),
+                          p=2)
+        return norm
+
+    def load_state_dict(self, state_dict):
+        super().load_state_dict(state_dict)
+        self.base_optimizer.param_groups = self.param_groups
+
+
+def disable_running_stats(model):
+
+    def _disable(module):
+        if isinstance(module, _BatchNorm):
+            module.backup_momentum = module.momentum
+            module.momentum = 0
+
+    model.apply(_disable)
+
+
+def enable_running_stats(model):
+
+    def _enable(module):
+        if isinstance(module, _BatchNorm) and hasattr(module,
+                                                      'backup_momentum'):
+            module.momentum = module.backup_momentum
+
+    model.apply(_enable)

attacks/AIM/src/gat/models/surrogate/__init__.py

@@ -0,0 +1,10 @@
+from .builder import build_surrogate, list_surrogates
+from .hooks import feat_col, midlayer_dict, register_collecter, register_collecter_cl
+from .tv import (densenet121, densenet169, inception_v3, resnet50, resnet152,
+                 swin_b, vgg16, vgg19, vit_b_16, vit_b_32)
+
+__all__ = [
+    'build_surrogate', 'list_surrogates', 'inception_v3', 'vgg16', 'vgg19',
+    'resnet50', 'resnet152', 'densenet121', 'densenet169', 'vit_b_16',
+    'vit_b_32', 'swin_b', 'midlayer_dict', 'register_collecter', 'feat_col', 'register_collecter_cl'
+]

attacks/AIM/src/gat/models/surrogate/builder.py

@@ -0,0 +1,11 @@
+from ...runtime.factory import Registry
+
+SURROGATE_REGISTRY = Registry('SURROGATE')
+
+
+def build_surrogate(_type: str, *args, **kwargs) -> object:
+    return SURROGATE_REGISTRY.get(_type)(*args, **kwargs)
+
+
+def list_surrogates():
+    return list(SURROGATE_REGISTRY._obj_map.keys())

attacks/AIM/src/gat/models/surrogate/hooks.py

@@ -0,0 +1,59 @@
+from contextlib import contextmanager
+from typing import List, Union
+
+import torch
+
+midlayer_dict = {
+    'vgg16': 'features.16',
+    'vgg19': 'features.18',
+    'resnet152': 'layer1',
+    'densenet169': 'features.denseblock2',
+    'inception_v3': 'Mixed_6c',
+    'resnet50': 'conv1',  # conv1, layer1, layer2, layer3, layer4
+    'resnet50_cl': 'layer4',
+    'resnet32': 'conv_1_3x3',   # conv_1_3x3, stage_1, stage_2, stage_3
+    'resnet32_cl': 'conv_1_3x3',
+}
+
+
+def register_collecter(m: torch.nn.Module, layer: str, feat_collecter: List):
+
+    def _hook(m, i, o):
+        feat_collecter.append(o)
+
+    _handler = m.get_submodule(layer).register_forward_hook(_hook)  #m.convnets[0].get_submodule(layer).register_forward_hook(_hook)
+    return _handler, feat_collecter
+
+def register_collecter_cl(m: torch.nn.Module, layer: str, feat_collecter: List, cl_methods: str):
+
+    def _hook(m, i, o):
+        feat_collecter.append(o)
+
+    if cl_methods == 'icarl' or cl_methods == 'finetune' or cl_methods == 'wa' or cl_methods == 'replay' or cl_methods == 'podnet' or cl_methods == 'bic':
+        _handler = m.convnet.get_submodule(layer).register_forward_hook(_hook)    # For ewc, icarl methods
+    elif cl_methods == 'foster' or cl_methods == 'der':
+        _handler = m.convnets[0].get_submodule(layer).register_forward_hook(_hook)  # For foster and der methods
+    elif cl_methods == 'memo':
+        _handler = m.TaskAgnosticExtractor.get_submodule(layer).register_forward_hook(_hook)  # For foster and der methods
+    return _handler, feat_collecter
+
+
+@contextmanager
+def feat_col(m: Union[torch.nn.Module, List[torch.nn.Module]],
+             layer: Union[str, List[str]]):
+    if isinstance(m, torch.nn.Module):
+        m = [m]
+    if isinstance(layer, str):
+        layer = [layer]
+    assert len(m) == len(layer)
+    handlers = []
+    feat_collecter = []
+    for _m, _layer in zip(m, layer):
+        handler, feat_collecter = register_collecter(_m, _layer,
+                                                     feat_collecter)
+        handlers.append(handler)
+    yield feat_collecter
+    for handler in handlers:
+        handler.remove()
+    feat_collecter.clear()
+    del feat_collecter

attacks/AIM/src/gat/models/surrogate/tv.py

@@ -0,0 +1,206 @@
+from torchvision import models as tv_models
+
+from .builder import SURROGATE_REGISTRY
+
+weights_dict = dict(
+    alexnet=tv_models.AlexNet_Weights.IMAGENET1K_V1,
+    googlenet=tv_models.GoogLeNet_Weights.IMAGENET1K_V1,
+    vgg16=tv_models.VGG16_Weights.IMAGENET1K_V1,
+    vgg19=tv_models.VGG19_Weights.IMAGENET1K_V1,
+    inception_v3=tv_models.Inception_V3_Weights.IMAGENET1K_V1,
+    resnet18=tv_models.ResNet18_Weights.IMAGENET1K_V1,
+    resnet34=tv_models.ResNet34_Weights.IMAGENET1K_V1,
+    resnet50=tv_models.ResNet50_Weights.IMAGENET1K_V1,
+    resnet152=tv_models.ResNet152_Weights.IMAGENET1K_V1,
+    wide_resnet50_2=tv_models.Wide_ResNet50_2_Weights.IMAGENET1K_V1,
+    wide_resnet101_2=tv_models.Wide_ResNet101_2_Weights.IMAGENET1K_V1,
+    densenet121=tv_models.DenseNet121_Weights.IMAGENET1K_V1,
+    densenet169=tv_models.DenseNet169_Weights.IMAGENET1K_V1,
+    mobilenet_v2=tv_models.MobileNet_V2_Weights.IMAGENET1K_V1,
+    mobilenet_v3_small=tv_models.MobileNet_V3_Small_Weights.IMAGENET1K_V1,
+    mobilenet_v3_large=tv_models.MobileNet_V3_Large_Weights.IMAGENET1K_V1,
+    squeezenet1_0=tv_models.SqueezeNet1_0_Weights.IMAGENET1K_V1,
+    squeezenet1_1=tv_models.SqueezeNet1_1_Weights.IMAGENET1K_V1,
+    shufflenet_v2_x0_5=tv_models.ShuffleNet_V2_X0_5_Weights.IMAGENET1K_V1,
+    shufflenet_v2_x1_0=tv_models.ShuffleNet_V2_X1_0_Weights.IMAGENET1K_V1,
+    shufflenet_v2_x1_5=tv_models.ShuffleNet_V2_X1_5_Weights.IMAGENET1K_V1,
+    shufflenet_v2_x2_0=tv_models.ShuffleNet_V2_X2_0_Weights.IMAGENET1K_V1,
+    efficientnet_b0=tv_models.EfficientNet_B0_Weights.IMAGENET1K_V1,
+    efficientnet_v2_s=tv_models.EfficientNet_V2_S_Weights.IMAGENET1K_V1,
+    efficientnet_v2_m=tv_models.EfficientNet_V2_M_Weights.IMAGENET1K_V1,
+    efficientnet_v2_l=tv_models.EfficientNet_V2_L_Weights.IMAGENET1K_V1,
+    vit_b_16=tv_models.ViT_B_16_Weights.IMAGENET1K_V1,
+    vit_b_32=tv_models.ViT_B_32_Weights.IMAGENET1K_V1,
+    swin_b=tv_models.Swin_B_Weights.IMAGENET1K_V1)
+
+
+@SURROGATE_REGISTRY.register()
+def alexnet(pretrain=True):
+    return tv_models.alexnet(
+        weights=weights_dict['alexnet'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def googlenet(pretrain=True):
+    return tv_models.googlenet(
+        weights=weights_dict['googlenet'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def vgg16(pretrain=True):
+    return tv_models.vgg16(weights=weights_dict['vgg16'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def vgg19(pretrain=True):
+    return tv_models.vgg19(weights=weights_dict['vgg19'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def inception_v3(pretrain=True):
+    return tv_models.inception_v3(
+        weights=weights_dict['inception_v3'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def resnet18(pretrain=True):
+    return tv_models.resnet18(
+        weights=weights_dict['resnet18'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def resnet34(pretrain=True):
+    return tv_models.resnet34(
+        weights=weights_dict['resnet34'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def resnet50(pretrain=True):
+    return tv_models.resnet50(
+        weights=weights_dict['resnet50'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def resnet152(pretrain=True):
+    return tv_models.resnet152(
+        weights=weights_dict['resnet152'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def wide_resnet50_2(pretrain=True):
+    return tv_models.wide_resnet50_2(
+        weights=weights_dict['wide_resnet50_2'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def wide_resnet101_2(pretrain=True):
+    return tv_models.wide_resnet101_2(
+        weights=weights_dict['wide_resnet101_2'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def densenet121(pretrain=True):
+    return tv_models.densenet121(
+        weights=weights_dict['densenet121'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def densenet169(pretrain=True):
+    return tv_models.densenet169(
+        weights=weights_dict['densenet169'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def mobilenet_v2(pretrain=True):
+    return tv_models.mobilenet_v2(
+        weights=weights_dict['mobilenet_v2'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def mobilenet_v3_small(pretrain=True):
+    return tv_models.mobilenet_v3_small(
+        weights=weights_dict['mobilenet_v3_small'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def mobilenet_v3_large(pretrain=True):
+    return tv_models.mobilenet_v3_large(
+        weights=weights_dict['mobilenet_v3_large'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def squeezenet1_0(pretrain=True):
+    return tv_models.squeezenet1_0(
+        weights=weights_dict['squeezenet1_0'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def squeezenet1_1(pretrain=True):
+    return tv_models.squeezenet1_1(
+        weights=weights_dict['squeezenet1_1'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def shufflenet_v2_x0_5(pretrain=True):
+    return tv_models.shufflenet_v2_x0_5(
+        weights=weights_dict['shufflenet_v2_x0_5'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def shufflenet_v2_x1_0(pretrain=True):
+    return tv_models.shufflenet_v2_x1_0(
+        weights=weights_dict['shufflenet_v2_x1_0'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def shufflenet_v2_x1_5(pretrain=True):
+    return tv_models.shufflenet_v2_x1_5(
+        weights=weights_dict['shufflenet_v2_x1_5'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def shufflenet_v2_x2_0(pretrain=True):
+    return tv_models.shufflenet_v2_x2_0(
+        weights=weights_dict['shufflenet_v2_x2_0'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def efficientnet_b0(pretrain=True):
+    return tv_models.efficientnet_b0(
+        weights=weights_dict['efficientnet_b0'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def efficientnet_v2_s(pretrain=True):
+    return tv_models.efficientnet_v2_s(
+        weights=weights_dict['efficientnet_v2_s'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def efficientnet_v2_m(pretrain=True):
+    return tv_models.efficientnet_v2_m(
+        weights=weights_dict['efficientnet_v2_m'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def efficientnet_v2_l(pretrain=True):
+    return tv_models.efficientnet_v2_l(
+        weights=weights_dict['efficientnet_v2_l'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def vit_b_16(pretrain=True):
+    return tv_models.vit_b_16(
+        weights=weights_dict['vit_b_16'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def vit_b_32(pretrain=True):
+    return tv_models.vit_b_32(
+        weights=weights_dict['vit_b_32'] if pretrain else None)
+
+
+@SURROGATE_REGISTRY.register()
+def swin_b(pretrain=True):
+    return tv_models.swin_b(
+        weights=weights_dict['swin_b'] if pretrain else None)

attacks/AIM/src/gat/runtime/__init__.py

@@ -0,0 +1,4 @@
+from .meter import AverageMeter, calc_cls_accuracy
+from .utils import fix_random, randid
+
+__all__ = ['fix_random', 'randid', 'AverageMeter', 'calc_cls_accuracy']

attacks/AIM/src/gat/runtime/api/aim_attack.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+import argparse
+import io
+from pathlib import Path
+from typing import Union
+
+import torchvision
+import uvicorn
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import StreamingResponse
+from PIL import Image
+
+from ...datasets.transforms import resize_256_224, to_color, to_pil, to_ts
+from ...models.attack import AIMAttack
+
+
+def init_attack(ckpt: Union[str, Path] = None):
+    attack = AIMAttack(device='cpu')
+    if ckpt:
+        attack.load_ckpt(ckpt)
+    attack.set_mode('eval')
+    attack_preproc = torchvision.transforms.Compose(resize_256_224() +
+                                                    to_color() + to_ts())
+    return attack, attack_preproc
+
+
+def main():
+    parser = argparse.ArgumentParser(description='AIM Attack API')
+    parser.add_argument('--host', type=str, default='0.0.0.0', help='host')
+    parser.add_argument('--port', type=int, default=8000, help='port')
+    parser.add_argument('--ckpt',
+                        type=str,
+                        default=None,
+                        help='path to the checkpoint')
+    args = parser.parse_args()
+
+    attack, attack_preproc = init_attack(args.ckpt)
+    app = FastAPI()
+
+    @app.get('/attack/aim/')
+    async def aim_attack(x_nat: UploadFile = File(...),
+                         x_guid: UploadFile = File(...)):
+        io_x_nat = await x_nat.read()
+        io_x_guid = await x_guid.read()
+        pil_x_nat = Image.open(io.BytesIO(io_x_nat))
+        pil_x_guid = Image.open(io.BytesIO(io_x_guid))
+        ts_x_nat = attack_preproc(pil_x_nat).unsqueeze(0)
+        ts_x_guid = attack_preproc(pil_x_guid).unsqueeze(0)
+        ts_x_adv = attack(ts_x_nat, ts_x_guid)
+        pil_x_adv = torchvision.transforms.Compose(to_pil())(ts_x_adv[0])
+
+        img_byte_array = io.BytesIO()
+        pil_x_adv.save(img_byte_array, format='PNG')
+        img_byte_array.seek(0)
+
+        return StreamingResponse(img_byte_array, media_type='image/png')
+
+    uvicorn.run(app, host=args.host, port=args.port)

attacks/AIM/src/gat/runtime/factory.py

@@ -0,0 +1,97 @@
+# Code is copied from
+# github.com/facebookresearch/fvcore:common/registry.py@242366
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# pyre-strict
+# pyre-ignore-all-errors[2,3]
+from typing import Any, Dict, Iterable, Iterator, Tuple
+
+from tabulate import tabulate
+
+
+class Registry(Iterable[Tuple[str, Any]]):
+    """
+    The registry that provides name -> object mapping, to support third-party
+    users' custom modules.
+
+    To create a registry (e.g. a backbone registry):
+
+    .. code-block:: python
+
+        BACKBONE_REGISTRY = Registry('BACKBONE')
+
+    To register an object:
+
+    .. code-block:: python
+
+        @BACKBONE_REGISTRY.register()
+        class MyBackbone():
+            ...
+
+    Or:
+
+    .. code-block:: python
+
+        BACKBONE_REGISTRY.register(MyBackbone)
+    """
+
+    def __init__(self, name: str) -> None:
+        """
+        Args:
+            name (str): the name of this registry
+        """
+        self._name: str = name
+        self._obj_map: Dict[str, Any] = {}
+
+    def _do_register(self, name: str, obj: Any) -> None:
+        assert (
+            name not in self._obj_map
+        ), "An object named '{}' was already registered in '{}' registry!".format(  # noqa: E501
+            name, self._name
+        )
+        self._obj_map[name] = obj
+
+    def register(self, obj: Any = None) -> Any:
+        """
+        Register the given object under the the name `obj.__name__`.
+        Can be used as either a decorator or not. See docstring of this class
+        for usage.
+        """
+        if obj is None:
+            # used as a decorator
+            def deco(func_or_class: Any) -> Any:
+                name = func_or_class.__name__
+                self._do_register(name, func_or_class)
+                return func_or_class
+
+            return deco
+
+        # used as a function call
+        name = obj.__name__
+        self._do_register(name, obj)
+
+    def get(self, name: str) -> Any:
+        ret = self._obj_map.get(name)
+        if ret is None:
+            raise KeyError(
+                "No object named '{}' found in '{}' registry!".format(
+                    name, self._name
+                )
+            )
+        return ret
+
+    def __contains__(self, name: str) -> bool:
+        return name in self._obj_map
+
+    def __repr__(self) -> str:
+        table_headers = ['Names', 'Objects']
+        table = tabulate(
+            self._obj_map.items(), headers=table_headers, tablefmt='fancy_grid'
+        )
+        return 'Registry of {}:\n'.format(self._name) + table
+
+    def __iter__(self) -> Iterator[Tuple[str, Any]]:
+        return iter(self._obj_map.items())
+
+    # pyre-fixme[4]: Attribute must be annotated.
+    __str__ = __repr__

attacks/AIM/src/gat/runtime/meter.py

@@ -0,0 +1,27 @@
+class AverageMeter:
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def calc_cls_accuracy(output, target, topk=(1,)):
+    maxk = min(max(topk), output.size()[1])
+    batch_size = target.size(0)
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.reshape(1, -1).expand_as(pred))
+    return [
+        correct[: min(k, maxk)].reshape(-1).float().sum(0) * 100.0 / batch_size
+        for k in topk
+    ]

attacks/AIM/src/gat/runtime/utils.py

@@ -0,0 +1,18 @@
+import random
+
+import numpy as np
+import torch
+
+
+def fix_random(seed: int = 0) -> None:
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def randid(k: int = 4):
+    charset = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'
+    return ''.join(random.choices(charset, k=k))

attacks/AIM/tests/test_datasets/test_datasets.py

@@ -0,0 +1,33 @@
+import os
+import unittest
+from pathlib import Path
+
+from gat.datasets import build_dataset
+
+in1k_data_root = '/root/workspace/proj/transfer-at/data/in_1k'
+in1k_data_root = os.environ.get('DATA_ROOT',
+                                Path(__file__).parents[2] / 'data' / 'in_1k')
+
+
+class TestImageNet(unittest.TestCase):
+
+    def test_in1k(self):
+        ds = build_dataset(
+            'imagenet',
+            data_root=in1k_data_root,
+            is_train=True,
+        )
+        self.assertEqual(len(ds), 1281167)
+
+    def test_in1k_filter(self):
+        ds = build_dataset(
+            'imagenet',
+            data_root=in1k_data_root,
+            is_train=False,
+            filter_class=0,
+        )
+        self.assertEqual(len(ds), 50)
+
+
+if __name__ == '__main__':
+    unittest.main()

attacks/AIM/tests/test_datasets/test_transforms.py

@@ -0,0 +1,80 @@
+import unittest
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose
+
+from gat.datasets.transforms import (hflip, norm, resize_224, resize_256_224,
+                                     resize_512_448, to_color, to_pil, to_ts)
+
+
+class TestResize(unittest.TestCase):
+
+    def test_resize_256_224(self):
+        inputs = torch.rand(3, 256, 256)
+        outputs = Compose(resize_256_224())(inputs)
+        self.assertEqual(outputs.shape, (3, 224, 224))
+
+    def test_resize_512_448(self):
+        inputs = torch.rand(3, 512, 512)
+        outputs = Compose(resize_512_448())(inputs)
+        self.assertEqual(outputs.shape, (3, 448, 448))
+
+    def test_resize_224(self):
+        inputs = torch.rand(3, 224, 224)
+        outputs = Compose(resize_224())(inputs)
+        self.assertEqual(outputs.shape, (3, 224, 224))
+
+
+class TestAug(unittest.TestCase):
+
+    def test_probability_invalid(self):
+        with self.assertRaises(AssertionError):
+            hflip(-0.1)
+
+    def test_probability_valid(self):
+        for valid_p in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]:
+            result = hflip(valid_p)
+            self.assertEqual(result[0].p, valid_p)
+
+
+class TestTypeTransforms(unittest.TestCase):
+
+    def test_to_ts(self):
+        pil_image = Image.new('RGB', (224, 224))
+        outputs = Compose(to_ts())(pil_image)
+        self.assertIsInstance(outputs, torch.Tensor)
+
+    def test_to_pil(self):
+        inputs = torch.rand(3, 224, 224)
+        outputs = Compose(to_pil())(inputs)
+        self.assertIsInstance(outputs, Image.Image)
+
+    def test_to_color_grayscale(self):
+        inputs = Image.new('L', (224, 224))
+        outputs = Compose(to_color())(inputs)
+        self.assertEqual(outputs.mode, 'RGB')
+
+    def test_to_color_rgb(self):
+        inputs = Image.new('RGB', (224, 224))
+        outputs = Compose(to_color())(inputs)
+        self.assertEqual(outputs.mode, 'RGB')
+
+
+class TestNorm(unittest.TestCase):
+
+    def test_default(self):
+        self.assertEqual(norm()[0].mean, (0.485, 0.456, 0.406))
+        self.assertEqual(norm()[0].std, (0.229, 0.224, 0.225))
+
+    def test_imagenet(self):
+        self.assertEqual(norm('IMAGENET')[0].mean, (0.485, 0.456, 0.406))
+        self.assertEqual(norm('IMAGENET')[0].std, (0.229, 0.224, 0.225))
+
+    def test_invalid_ds(self):
+        with self.assertRaises(AttributeError):
+            norm('imagenets')
+
+
+if __name__ == '__main__':
+    unittest.main()

attacks/AIM/tests/test_models/test_attack.py

@@ -0,0 +1,110 @@
+import unittest
+
+import torch
+
+from gat.models.attack.aim_attack import AIMAttack
+from gat.models.attack.cda_attack import CDAAttack
+from gat.models.attack.generator.aim import AIMGenerator
+from gat.models.attack.generator.cda import CDAGenerator
+from gat.models.attack.loss.logits import ContrastiveLoss
+
+
+class TestGenerator(unittest.TestCase):
+
+    @torch.no_grad()
+    def test_cda(self):
+        generator = CDAGenerator()
+        inputs = torch.rand(1, 3, 224, 224)
+        outputs = generator(inputs)
+        self.assertEqual(outputs.shape, (1, 3, 224, 224))
+
+    @torch.no_grad()
+    def test_aim(self):
+        generator = AIMGenerator()
+        inputs = (torch.rand(1, 3, 224, 224), torch.rand(1, 3, 224, 224))
+        outputs = generator(*inputs)
+        self.assertEqual(outputs.shape, (1, 3, 224, 224))
+
+
+class TestContrastiveLoss(unittest.TestCase):
+
+    def setUp(self):
+        self.margin = 1.0
+        self.loss_fn = ContrastiveLoss(margin=self.margin)
+
+    @torch.no_grad()
+    def test_forward_shape(self):
+        anchors = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+        positives = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+        negatives = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
+
+        loss = self.loss_fn(anchors, negatives, positives)
+        self.assertEqual(loss.shape, torch.Size([]))
+
+    @torch.no_grad()
+    def test_loss_outputs(self):
+        anchors = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+        positives = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+        negatives = torch.tensor([[2.0, 0.0], [0.0, 2.0]])
+
+        loss = self.loss_fn(anchors, negatives, positives)
+        self.assertAlmostEqual(loss.item(), 0.25)
+
+    @torch.no_grad()
+    def test_non_zero_loss(self):
+        anchors = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+        positives = torch.tensor([[0.0, 1.0], [1.0, 0.0]])
+        negatives = torch.tensor([[2.0, 0.0], [0.0, 2.0]])
+
+        loss = self.loss_fn(anchors, negatives, positives)
+        self.assertAlmostEqual(loss.item(), 0.75)
+
+
+class TestCDAAttack(unittest.TestCase):
+
+    def setUp(self):
+        self.device = torch.device('cpu')
+        self.epsilon = 16. / 255.
+        self.attack = CDAAttack(self.device, self.epsilon)
+
+    @torch.no_grad()
+    def test_outputs_shape(self):
+        x_nat = torch.rand(1, 3, 224, 224)
+        x_adv = self.attack(x_nat)
+        self.assertEqual(x_adv.shape, (1, 3, 224, 224))
+
+    @torch.no_grad()
+    def test_outputs_bound(self):
+        x_nat = torch.rand(1, 3, 224, 224)
+        x_adv = self.attack(x_nat)
+        self.assertTrue((x_adv >= 0.0).all())
+        self.assertTrue((x_adv <= 1.0).all())
+        self.assertTrue((x_adv - x_nat).abs().max() <= self.epsilon)
+
+
+class TestAIMAttack(unittest.TestCase):
+
+    def setUp(self):
+        self.device = torch.device('cpu')
+        self.epsilon = 16. / 255.
+        self.attack = AIMAttack(self.device, self.epsilon)
+
+    @torch.no_grad()
+    def test_outputs_shape(self):
+        x_nat = torch.rand(1, 3, 224, 224)
+        x_guid = torch.rand(1, 3, 224, 224)
+        x_adv = self.attack(x_nat, x_guid)
+        self.assertEqual(x_adv.shape, (1, 3, 224, 224))
+
+    @torch.no_grad()
+    def test_outputs_bound(self):
+        x_nat = torch.rand(1, 3, 224, 224)
+        x_guid = torch.rand(1, 3, 224, 224)
+        x_adv = self.attack(x_nat, x_guid)
+        self.assertTrue((x_adv >= 0.0).all())
+        self.assertTrue((x_adv <= 1.0).all())
+        self.assertTrue((x_adv - x_nat).abs().max() <= self.epsilon)
+
+
+if __name__ == '__main__':
+    unittest.main()

attacks/AIM/tests/test_models/test_surrogate.py

@@ -0,0 +1,59 @@
+import unittest
+
+import torch
+
+from gat.models.surrogate import build_surrogate, list_surrogates
+from gat.models.surrogate.hooks import feat_col
+
+
+class TestTVModels(unittest.TestCase):
+
+    @torch.no_grad()
+    def test_outputs_shape(self):
+        inputs = torch.rand(1, 3, 224, 224)
+        for surrogate_id in list_surrogates():
+            if surrogate_id not in ['inception_v3']:
+                outputs = build_surrogate(surrogate_id, pretrain=False)(inputs)
+                self.assertEqual(outputs.shape, (1, 1000))
+
+
+class TestFeatCol(unittest.TestCase):
+
+    @torch.no_grad()
+    def test_feat_col(self):
+        testcases = [{
+            'surrogate_id': 'vgg16',
+            'feat_layer': 'features.16',
+            'input_shape': (1, 3, 224, 224),
+            'feat_shape': (1, 256, 28, 28)
+        }, {
+            'surrogate_id': 'vgg19',
+            'feat_layer': 'features.18',
+            'input_shape': (1, 3, 224, 224),
+            'feat_shape': (1, 256, 28, 28)
+        }, {
+            'surrogate_id': 'resnet152',
+            'feat_layer': 'layer2',
+            'input_shape': (1, 3, 224, 224),
+            'feat_shape': (1, 512, 28, 28)
+        }, {
+            'surrogate_id': 'densenet169',
+            'feat_layer': 'features.denseblock2',
+            'input_shape': (1, 3, 224, 224),
+            'feat_shape': (1, 512, 28, 28)
+        }]
+        for testcase in testcases:
+            model = build_surrogate(testcase['surrogate_id'], pretrain=False)
+            with feat_col(model, testcase['feat_layer']) as _feat_collecter:
+                inputs = torch.rand(testcase['input_shape'])
+                model(inputs)
+                model(inputs)
+                self.assertEqual(len(_feat_collecter), 2)
+                self.assertEqual(_feat_collecter.pop().shape,
+                                 testcase['feat_shape'])
+                self.assertEqual(_feat_collecter.pop().shape,
+                                 testcase['feat_shape'])
+
+
+if __name__ == '__main__':
+    unittest.main()