attack.py

import torch
import torch.nn as nn

import numpy as np

from utils import *

class Attack(object):
    """
    Base class for all attacks.
    """
    def __init__(self, attack, model_name, epsilon, targeted, random_start, norm, loss, device=None):
        """
        Initialize the hyperparameters

        Arguments:
            attack (str): the name of attack.
            model_name (str): the name of surrogate model for attack.
            epsilon (float): the perturbation budget.
            targeted (bool): targeted/untargeted attack.
            random_start (bool): whether using random initialization for delta.
            norm (str): the norm of perturbation, l2/linfty.
            loss (str): the loss function.
            device (torch.device): the device for data. If it is None, the device would be same as model
        """
        if norm not in ['l2', 'linfty']:
            raise Exception("Unsupported norm {}".format(norm))
        self.attack = attack
        self.model = model_name
        #self.model = self.load_ens_model()
        self.epsilon = epsilon
        self.targeted = targeted
        self.random_start = random_start
        self.norm = norm
        #self.device = next(self.model.models[0].parameters()).device if device is None else device
        self.device = next(self.model.parameters()).device if device is None else device
        self.loss = self.loss_function(loss)

    def load_model(self, model_name):
        """
        The model Loading stage, which should be overridden when surrogate model is customized (e.g., DSM, SE_TR, etc.)
        Prioritize the model in torchvision.models, then timm.models

        Arguments:
            model_name (str): the name of surrogate model in model_list in utils.py

        Returns:
            model (torch.nn.Module): the surrogate model wrapped by wrap_model in utils.py
        """
        if model_name in models.__dict__.keys():
            print('=> Loading model {} from torchvision.models'.format(model_name))
            model = models.__dict__[model_name](weights="IMAGENET1K_V1")
        elif model_name in timm.list_models():
            print('=> Loading model {} from timm.models'.format(model_name))
            model = timm.create_model(model_name, pretrained=True)
        else:
            raise ValueError('Model {} not supported'.format(model_name))
        return nn.DataParallel(wrap_model(model.eval().cuda()))
        return wrap_model(model.eval().cuda())
        #return wrap_model(model.eval().to(self.device))
        
    def load_ens_model(self):
        """
        The model Loading stage, which should be overridden when surrogate model is customized (e.g., DSM, SE_TR, etc.)
        Prioritize the model in torchvision.models, then timm.models

        Arguments:
            model_name (str): the name of surrogate model in model_list in utils.py

        Returns:
            model (torch.nn.Module): the surrogate model wrapped by wrap_model in utils.py
        """
        '''pretrain_name= ['resnet18', 'resnet101', 'resnext50_32x4d', 'densenet121', 'inception_v3', 
                      'inception_v4','vit_base_patch16_224', 'pit_b_224','visformer_small', 'swin_tiny_patch4_window7_224']'''
                      
        pretrain_name= ['resnet18', 'inception_v4','visformer_small', 'swin_tiny_patch4_window7_224']
        
        model_list = []
        for model_name in pretrain_name:
            if model_name in models.__dict__.keys():
                model = models.__dict__[model_name](weights="IMAGENET1K_V1")
            elif model_name in timm.list_models():
                model = timm.create_model(model_name, pretrained=True)
            model_list.append(nn.DataParallel(wrap_model(model.eval().cuda())))
        
        return EnsembleModel(model_list)

    def forward(self, data, label, **kwargs):
        """
        The general attack procedure

        Arguments:
            data (N, C, H, W): tensor for input images
            labels (N,): tensor for ground-truth labels if untargetd
            labels (2,N): tensor for [ground-truth, targeted labels] if targeted
        """
        if self.targeted:
            assert len(label) == 2
            label = label[1] # the second element is the targeted label tensor
        data = data.clone().detach().to(self.device)
        label = label.clone().detach().to(self.device)

        # Initialize adversarial perturbation
        delta = self.init_delta(data)

        momentum = 0
        for _ in range(self.epoch):
            # Obtain the output
            logits = self.get_logits(self.transform(data+delta, momentum=momentum))

            # Calculate the loss
            loss = self.get_loss(logits, label)

            # Calculate the gradients
            grad = self.get_grad(loss, delta)

            # Calculate the momentum
            momentum = self.get_momentum(grad, momentum)

            # Update adversarial perturbation
            delta = self.update_delta(delta, data, momentum, self.alpha)

        return delta.detach()

    def get_logits(self, x, **kwargs):
        """
        The inference stage, which should be overridden when the attack need to change the models (e.g., ensemble-model attack, ghost, etc.) or the input (e.g. DIM, SIM, etc.)
        """
        return self.model(x)

    def get_loss(self, logits, label):
        """
        The loss calculation, which should be overrideen when the attack change the loss calculation (e.g., ATA, etc.)
        """
        # Calculate the loss
        return -self.loss(logits, label) if self.targeted else self.loss(logits, label)


    def get_grad(self, loss, delta, **kwargs):
        """
        The gradient calculation, which should be overridden when the attack need to tune the gradient (e.g., TIM, variance tuning, enhanced momentum, etc.)
        """
        return torch.autograd.grad(loss, delta, retain_graph=False, create_graph=False)[0]

    def get_momentum(self, grad, momentum, **kwargs):
        """
        The momentum calculation
        """
        return momentum * self.decay + grad / (grad.abs().mean(dim=(1,2,3), keepdim=True))

    def init_delta(self, data, **kwargs):
        delta = torch.zeros_like(data).to(self.device)
        if self.random_start:
            if self.norm == 'linfty':
                delta.uniform_(-self.epsilon, self.epsilon)
            else:
                delta.normal_(-self.epsilon, self.epsilon)
                d_flat = delta.view(delta.size(0), -1)
                n = d_flat.norm(p=2, dim=10).view(delta.size(0), 1, 1, 1)
                r = torch.zeros_like(data).uniform_(0,1).to(self.device)
                delta *= r/n*self.epsilon
            delta = clamp(delta, img_min-data, img_max-data)
        delta.requires_grad = True
        return delta

    def update_delta(self, delta, data, grad, alpha, **kwargs):
        if self.norm == 'linfty':
            delta = torch.clamp(delta + alpha * grad.sign(), -self.epsilon, self.epsilon)
        else:
            grad_norm = torch.norm(grad.view(grad.size(0), -1), dim=1).view(-1, 1, 1, 1)
            scaled_grad = grad / (grad_norm + 1e-20)
            delta = (delta + scaled_grad * alpha).view(delta.size(0), -1).renorm(p=2, dim=0, maxnorm=self.epsilon).view_as(delta)
        # delta = torch.clamp(delta, img_min-data, img_max-data)
        return delta

    def loss_function(self, loss):
        """
        Get the loss function
        """
        if loss == 'crossentropy':
            return nn.CrossEntropyLoss()
        else:
            raise Exception("Unsupported loss {}".format(loss))

    def transform(self, data, **kwargs):
        return data

    def __call__(self, *input, **kwargs):
        self.model.eval()
        return self.forward(*input, **kwargs)