HFC_Attack.py

import argparse
import os
from tqdm import tqdm
import numpy as np

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
from torch.autograd import Variable

from datasets import get_dataloader
from utils import *
from network import *
from saver import Saver

from sklearn.metrics import accuracy_score

'''
Different attacks generated by gen_adv_samples are controled by args:
These attacks are implenmented according to https://adversarial-attacks-pytorch.readthedocs.io/en/latest/

BIM: By default
FGSM: setting steps = 1
PGD: setting rand_start = True
MIM: setting momentum = 0.5
DIM: setting diverse_prob = 0.5
TIM: setting TI_kernel = torch.from_numpy(kernel_generation()).cuda()
CW: setting cls_loss_fn = fn_cw_loss
'''

def gen_adv_samples(model, loss_fn, x, y, attack_name, eps=1 / 255, steps=10, \
        rand_start=False, eps_step=0, momentum=0, TI_kernel=None):

    # Attack types
    if attack_name == 'MI_FGSM_Linf': momentum = 0.5
    diverse_prob = 0.5 if attack_name == 'DI_FGSM_Linf' else 0
    if attack_name == 'TI_FGSM_Linf':
        TI_kernel = torch.from_numpy(kernel_generation()).cuda()
    elif attack_name == 'PGD_Linf':
        rand_start = True

    x_adv = x.clone()
    if rand_start:
        x_adv += torch.rand(x_adv.shape).cuda() * eps
    x_adv = Variable(x_adv.cuda(), requires_grad=True)
    x_adv.data = torch.clamp(x_adv, 0, 1)
    x = x.cuda()

    g = torch.zeros_like(x)
    for t in range(steps):

        x_adv.requires_grad = True
        out = model.feature_list(input_diversity(x_adv, diverse_prob))
        loss = loss_fn(out, y)
        loss.backward(retain_graph=True)
        
        grad = x_adv.grad.data
        if TI_kernel is not None:
            grad = F.conv2d(grad, TI_kernel, stride=1, padding=3, groups=3)
        g = momentum * g + normalize_by_pnorm(grad)

        x_adv.data = x_adv.data - eps_step * g.sign()
        x_adv.data = linfball_proj(x, eps, x_adv)
        x_adv.data = torch.clamp(x_adv, 0, 1)
        x_adv.grad = None

    return x_adv 

def run(args, attack_methods):
    num_component_gmm = args.num_component

    assert(type(attack_methods) == str)
    is_targeted = False
    saver = Saver(args.arch, dataset=args.dataset)
    logging.info(f'Attacking {args.arch}')

    test_loader = get_dataloader(dataset=args.dataset, arch=args.arch, mode='test',\
        batch_size=args.batch_size, num_workers=4, \
        num_fold=args.num_fold, targeted=is_targeted)
    num_classes = test_loader.dataset.num_classes

    if 'vgg16' in args.arch:
        src_model = infer_Cls_Net_vgg(num_classes)
    elif 'resnet50' in args.arch:
        src_model = infer_Cls_Net_resnet(num_classes)
    elif 'resnet3d' in args.arch:
        src_model = infer_Cls_Net_resnet3d(num_classes)
    else:
        raise NotImplementedError

    src_model = saver.load_model(src_model, args.arch)
    src_model.eval()
    src_model = src_model.cuda()

    num_layers = src_model.num_feature
    num_cnn_layers = src_model.num_cnn

    # Metric Savers
    metric_counter = dict()
    metric_counter = dict()
    gt = np.array([])
    pred = np.array([])
    metric_counter = dict()
    metric_counter['gt'] = gt
    metric_counter['pred'] = pred
    metric_counter['data'] = list()
    metric_counter['mse'] = list()
    metric_counter['mse_raw'] = list()

    cls_loss_fn = torch.nn.CrossEntropyLoss()

    def fn_cw_loss(logits, target, get_scaler=True):
        one_hot = torch.zeros_like(logits)
        one_hot = one_hot.scatter(1, target.view(-1,1), 1)
        target_logits = (one_hot * logits).sum(-1)
        remaining_logits_max = ((1 - one_hot) * logits).max(-1)[0]
        if get_scaler:
            cw_loss = torch.clamp(remaining_logits_max - target_logits + 100, min=0.).mean()
        else:
            cw_loss = torch.clamp(remaining_logits_max - target_logits + 100, min=0.)
        return cw_loss
    if 'CW' in attack_methods[0]:
        cls_loss_fn = fn_cw_loss
        print("Use CW Attack Loss Function !")

    def loss_fn_(input, item):
        
        logits, feature = input
        target = item[0]
        cls_loss = cls_loss_fn(logits, target)
        final_loss = cls_loss
        mean_list, std_list, weights = item[1]

        # Compute HFC loss
        for id_layer in range(num_cnn_layers):
            # if id_layer < num_cnn_layers - 1: continue
            temp_mean = mean_list[id_layer]
            std = std_list[id_layer]
            weight = weights[id_layer].log()
            
            # Skip the layers with too large weights with big HFC losses
            shape = feature[id_layer].shape
            if shape[1] > 512: continue

            # Compute mean and prob to n_conponent
            mean_feature = feature[id_layer].view(feature[id_layer].size(0), \
                feature[id_layer].size(1), -1).mean(-1)
            num_component = temp_mean.shape[0]
            
            with torch.no_grad():
                select_featue = mean_feature.unsqueeze(1).repeat((1,num_component,1))
                # mean : [batch, n_comp, n_fea] == zero_f
                zero_f = select_featue - temp_mean
                scores = torch.zeros([zero_f.shape[0], zero_f.shape[1]]).float().cuda()
                for id_comp in range(num_component):
                    id_mean = zero_f[:,id_comp,:]
                    hfc_score = torch.mm(id_mean, std[id_comp]) * torch.mm(id_mean, std[id_comp])
                    score = -0.5 * hfc_score.sum(-1)
                    score = score + std[id_comp].diag().log().sum()
                    scores[:, id_comp] = score
                scores += weight
                # Find the nearest component
                selected_comp = scores.argmax(-1)
            
            zero_f = mean_feature - temp_mean[selected_comp]
            tmp =  torch.bmm(zero_f.unsqueeze(1), std[selected_comp]).squeeze()
            hfc_loss = (tmp * tmp).mean(-1).mean()
            final_loss =  args.lamda * hfc_loss + final_loss
        # print('bce {:.3f} hfc_loss {:.3f}'.format(cls_loss, hfc_loss))
        # print('bce {:.3f} final {:.3f}'.format(cls_loss, final_loss))
        return  final_loss
    
    loss_fn = loss_fn_

    # Define attack parameters
    name = attack_methods
    splits = name.split('_')
    attack_name = ('_').join(splits[:-1])
    epsilon = float(name.split('_')[-1]) / 256
    print(f'Attack by Constrain: {256 * epsilon}')


    if args.dataset == 'APTOS' and 'vgg16' in args.arch:
        eps_step = 2 / 256 / 20
    elif args.dataset == 'APTOS' and 'resnet50' in args.arch:
        eps_step = 0.5 / 256 / 20
    elif args.dataset == 'Brain' and 'resnet3d' in args.arch:
        eps_step = 0.5 / 256 / 20
    elif args.dataset == 'CXR' and 'vgg16' in args.arch:
        eps_step = 2 / 256 / 50
    elif args.dataset == 'CXR' and 'resnet50' in args.arch:
        eps_step = 0.5 / 256 / 50
    elif args.dataset == 'Cifar' and 'resnet50' in args.arch:
        eps_step = 0.5 / 256 / 10
    else:
        raise NotImplementedError

    if attack_name == 'FGSM_Linf':
        steps = 1
        eps_step = epsilon
    else:
        # We choose fixed eps_step attack for HFC Attack, and computes steps accoordingly.
        steps = 2 * int(epsilon / eps_step)
        pass
    print('Number of steps: {}'.format(steps))

    clean_feature_dict = dict()
    adv_feature_dict = dict()
    for i in range(2):
        clean_feature_dict[i] = list()
        adv_feature_dict[i] = list()

    total_feature_list = {item:[] for item in range(num_layers)}
    total_pred_list = list()
    clean_images = list() 
    
    root_gmm = os.path.join(os.getcwd(), f'runs_{args.dataset}', args.arch, \
        f'GMM_{num_component_gmm}')
    print(f'load GMM from {root_gmm}')
    
    # Start to HFC attack
    for id_class in range(2):
        test_loader = get_dataloader(dataset=args.dataset, arch=args.arch, mode='adv_test',\
            batch_size=args.batch_size, num_workers=4, \
            num_fold=args.num_fold, targeted=is_targeted, rand_pairs='specific', target_class=id_class)

        mean_list, pre_chol_list, weight_list, scores = [], [], [], []
        for id_layer in range(num_cnn_layers):
            gmm_model = np.load(os.path.join(root_gmm, f'Layer_{id_layer}_class_{id_class}.npz'))
            # Here precision is a metric with shape [n_components, n_featues, n_featues]
            # gmm.means_ shape [n_components, n_features]
            # precision[i] = covirence[i].inverse() (ni ju zhen)
            # precision cholesky is the cholesky decomposition of precision (for better computing only)
            mean_list.append(torch.from_numpy(gmm_model['means_']).cuda().float())
            pre_chol_list.append(torch.from_numpy(gmm_model['precision_cholesky_']).cuda().float())
            weight_list.append(torch.from_numpy(gmm_model['weights_']).cuda().float())

        for i, (images, target) in enumerate(tqdm(test_loader, desc=f'Class {id_class}')):
            # if i > 2: break
            if True:
                images = images.cuda()
                clean_images.append(images.cpu().numpy())
                target = target.cuda()
                # false_images = false_images.cuda()
            
            adv_images = gen_adv_samples(src_model, loss_fn, images, [target, \
                    [mean_list, pre_chol_list, weight_list]], attack_name,\
                steps=steps, eps=epsilon, eps_step=eps_step)

            output = src_model(adv_images).argmax(dim=1).detach().cpu().numpy()
            metric_counter['data'].append(adv_images.detach().cpu().numpy())
            gt_concatnate = target.detach().cpu().numpy()
            metric_counter['gt'] = \
                np.concatenate([metric_counter['gt'], gt_concatnate], axis=0)
            metric_counter['pred'] = \
                np.concatenate([metric_counter['pred'], output], axis=0)
    
            
            adv_feature_list = src_model.feature_list(adv_images)[1]
            for id_layer, item in enumerate(adv_feature_list):
                temp = item.view(item.shape[0], item.shape[1], -1).mean(-1).detach().cpu().numpy()
                total_feature_list[id_layer].append(temp)
            total_pred_list.append(output)

    # calculate metrics
    metric_counter['acc'] = accuracy_score(\
        metric_counter['gt'], \
        metric_counter['pred'], normalize=True)
    print("Acc {:.3f} Using {}".format(\
        metric_counter['acc'], attack_name))
    
    from eval import Feature_Detector
    detector = Feature_Detector(temp_dir=\
        os.path.join(os.getcwd(), f'runs_{args.dataset}', args.arch, args.detector),\
            num_classes=num_classes, num_layers=num_layers, num_cnn_layers=num_cnn_layers)

    save_imgs = np.concatenate(metric_counter['data'], axis=0)
    np.save(os.path.join(f'runs_{args.dataset}', args.arch, args.attack +'_HFC.npy'), save_imgs)
    adv_data = torch.from_numpy(save_imgs).cuda().float()
    bingo = metric_counter['gt'] == metric_counter['pred']
    correct_image = adv_data[bingo]
    labels = metric_counter['gt'][bingo]
    clean_images = np.concatenate(clean_images, axis=0)
    clean_images_correct = clean_images[bingo]
    if labels.shape[0] % 10 == 1:
        correct_image = correct_image[:10 * (labels.shape[0]//10)]
        labels = labels[:10 * (labels.shape[0]//10)]

    diff = np.abs(correct_image.detach().cpu().numpy() - clean_images_correct) * 256
    diff = diff.reshape(diff.shape[0], -1)
    print(f'L inf max {diff.max()}')

    detector.eval_patch(correct_image, src_model, labels)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Training Code')
    parser.add_argument('-a', '--arch', metavar='ARCH', default='vgg16')
    parser.add_argument('-b', '--batch-size', default=4, type=int,
                        metavar='N',
                        help='mini-batch size (default: 256), this is the total '
                            'batch size of all GPUs on the current node when '
                            'using Data Parallel or Distributed Data Parallel')
    parser.add_argument('-f', '--num_fold', default=0, type=int,
                        help='Fold Number')
    parser.add_argument('-y', '--num_component', default=1, type=int,
                        help='num_component')
    parser.add_argument('--dataset', default='APTOS', type=str,
                        help='Fold Number')
    parser.add_argument('--attack', default='I_FGSM_Linf_1', type=str,
                        help='Type of the adversarial attack') 
    parser.add_argument('--detector', default='I_FGSM_Linf_1', type=str,
                        help='Choose detector trained by which adversarial attack')  
    parser.add_argument('--lamda', default=1, type=float,
                        help='lamda') 
    parser.add_argument('--get_feature', default=0, type=float,
                        help='get feature layer index')                   
    parser.add_argument('--multiprocessing-distributed', action='store_true',
                        help='Use multi-processing distributed training to launch '
                            'N processes per node, which has N GPUs. This is the '
                            'fastest way to use PyTorch for either single node or '
                            'multi node data parallel training')
    args = parser.parse_args()

    print(f"Run Attacking {args.arch} using {args.attack}")
    run(args, args.attack)