attacks/Gaker/GAKERAttack.py

import os.path
import torchvision
import torch
from foolbox.attacks.base import *
from foolbox.attacks.gradient_descent_base import *
from tqdm import tqdm
import pandas as pd
from attacks.Gaker.Generator.Generator import Generator
from attacks.Gaker.Generator.train import CustomResnet50, CustomDenseNet121
from attacks.AIM.src.gat.models.surrogate import midlayer_dict, register_collecter_cl
from attacks.attack_config import SustainableAttack
from utils.plot import plot_asr_per_target, save_grad_cam
import logging
import foolbox as fb
from foolbox import PyTorchModel
import numpy as np
from utils import factory
from utils.data_manager import get_dataloader
from attacks.Gaker.utils_.gaussian_smoothing import get_gaussian_kernel

class Gaker(SustainableAttack):
    def __init__(self, args, device='cuda'):
        super().__init__(args, device)
        self.device = device
        self.args = args
        self.surrogate_model = None
        self.surrogate_model_name = 'resnet50'
        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.layer = midlayer_dict[self.surrogate_model_name]
            self.feat_collecter_handler, self.feat_collecter = register_collecter_cl(self.surrogate_model,
                                                                                     self.layer,
                                                                                     self.feat_collecter,
                                                                                     self.args["model_name"])
            self.feature_channel = 2048
        else:
            if self.surrogate_model_name == "resnet50":
                original_model = torchvision.models.resnet50(pretrained=True)
                self.feature_extraction = CustomResnet50(original_model)
                self.feature_extraction = self.feature_extraction.eval().to(self.device)
                self.feature_channel = 2048
            elif self.surrogate_model_name == "densenet121":
                original_model = torchvision.models.densenet121(pretrained=True)
                self.feature_extraction = CustomDenseNet121(original_model)
                self.feature_extraction = self.feature_extraction.eval().to(self.device)
                self.feature_channel = 1024
            elif self.surrogate_model_name == "vgg19bn":
                vgg19bn = torchvision.models.vgg19(pretrained=True).eval().to(self.device)
                self.feature_channel = 4096
                global hook_output
                hook_output = None
                def hook(module, input, output):
                    global hook_output
                    hook_output = output
                handle = vgg19bn.classifier[5].register_forward_hook(hook)

        self.adv_generator = Generator(num_target=10, ch=32, ch_mult=[1, 2, 3, 4],num_res_blocks=1,feature_channel_num=self.feature_channel).to(device)
        self.lr = 0.001  # 1e-4
        self.betas = (0.5, 0.999)
        self.num_epoch = 100
        self.optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, self.adv_generator.parameters()), lr=self.lr, weight_decay=5e-5)#GAN的优化器
        self.kernel = get_gaussian_kernel(kernel_size=3, pad=2, sigma=1).to(device)

        self.eps = 32 / 255
        self.ran_best = 'random'
        self.prefix = f'{int(self.eps*255)}'
        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.eval_batch_size = 128
        self.plot_gradcam = False

    def train_generator(self):
        self.file_path = os.path.join(self.save_path, f'{self.prefix}.pth')
        if os.path.exists(self.file_path):
            self.adv_generator.load_state_dict(torch.load(self.file_path, map_location=self.device))
            self.adv_generator.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)

            for epoch in range(1, self.num_epoch + 1):
                iteration = 0
                laoder_tqdm = tqdm(self.loader, total=len(self.loader), desc=f'Epoch {epoch}')
                for i, (_, x, y) in enumerate(laoder_tqdm):
                    x_f = x[y != self.target_class].to(self.device)
                    y_f = y[y != self.target_class].to(self.device)
                    del x, y

                    if len(x_f) > len(target_images):
                        x_f = x_f[:len(target_images)]
                        y_f = y_f[:len(target_images)]
                    else:
                        target_images = target_images[:len(x_f)]
                        target_labels = target_labels[:len(x_f)]

                    with torch.no_grad():
                        if self.surrogate_model_name == "resnet50" or self.surrogate_model_name == "densenet121":
                            target_fea = self.feature_extraction(self.norm(target_images)).squeeze()

                    output_to_mix = target_fea
                    target_feature = []
                    for i in range(target_images.shape[0]):
                        target_feature.append(target_fea[i])

                    target_feature = torch.tensor(
                        np.array([item.cpu().detach().numpy() for item in target_feature])).to(self.device)
                    mask = torch.ne(y_f, target_labels).long().to(self.device)
                    perturbated_imgs = self.kernel(self.adv_generator(x_f, mix=output_to_mix))

                    adv = torch.min(torch.max(perturbated_imgs, x_f - self.eps), x_f + self.eps)
                    adv = torch.clamp(adv, 0.0, 1.0)

                    if self.surrogate_model_name == "resnet50" or self.surrogate_model_name == "densenet121":
                        adv_feature = self.feature_extraction(self.norm(adv))

                    if self.surrogate_model_name == "resnet50" or self.surrogate_model_name == "densenet121" or \
                            self.surrogate_model_name == "vgg19bn":
                        adv_feature = adv_feature.squeeze()
                        loss = 1 - torch.cosine_similarity(adv_feature, target_feature, dim=1)

                    loss = mask * loss

                    noise = adv - x_f
                    if self.surrogate_model_name == "resnet50" or self.surrogate_model_name == "densenet121":
                        noise_feature = self.feature_extraction(self.norm(noise)).squeeze()
                        loss_noise = 1 - torch.cosine_similarity(noise_feature, target_feature, dim=1)

                    loss_noise = mask * loss_noise * 0.5
                    loss = loss + loss_noise

                    loss = (loss.sum()) / x_f.shape[0]

                    loss.backward()
                    self.optimizer.step()

                    iteration += 1

                    del x_f, y_f, noise, adv, adv_feature, loss, loss_noise, perturbated_imgs, mask
                    torch.cuda.empty_cache()
            torch.save(self.adv_generator.state_dict(), self.file_path)

    def run_test(self):
        # Load Batch Data
        self.adv_generator.eval()
        self.loader = get_dataloader(self.data_manager, batch_size=self.eval_batch_size,
                                start_class=0, end_class=10,
                                train=False, shuffle=True, num_workers=0)
        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)))

            target_imgs = imgs[labels == self.target_class]
            target_labels = labels[labels == self.target_class]
            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 = self.adv_gen(imgs.raw, target_image.raw.repeat(len(imgs), 1, 1, 1))
                is_adv = is_adversarial(ep.astensor(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.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()

        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 adv_gen(self, images, target_images):
        if self.ran_best == 'random':
            target_feature = self.feature_extraction(self.norm(target_images))
            output_to_mix = target_feature.squeeze()
        elif self.ran_best == 'best':
            print('not used')
        else:
            print('please choose random or best')

        perturbated_imgs = self.adv_generator(images, mix=output_to_mix)

        perturbated_imgs = self.kernel(perturbated_imgs)

        adv = torch.min(torch.max(perturbated_imgs, images - self.eps), images + self.eps)
        adv = torch.clamp(adv, 0, 1.0)

        return adv

    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) -> "Gaker":
        ...