gen_base_adv.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


'''
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(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):
    assert(type(attack_methods) == str)
    is_targeted = False
    saver = Saver(args.arch, dataset=args.dataset)
    logging.info(f'Attacking {args.arch}')

    # Load datasets
    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, rand_pairs='targeted_attack')

    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
    
    # Load model
    src_model = saver.load_model(src_model, args.arch)
    src_model.eval()
    src_model = src_model.cuda()

    # 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.)
        # print('wtf')
        return cw_loss
    
    # Define loss functions
    if 'CW' in attack_methods:
        cls_loss_fn = fn_cw_loss
        print("Use CW Attack Loss Function !")

    def loss_fn_warp(feature, item):
        
        logits = feature
        target = item[0]
        cls_loss = cls_loss_fn(logits, target)
        final_loss = cls_loss
        return  final_loss

    loss_fn = loss_fn_warp

    # 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}')

    # Steps : Attack steps for each image
    if attack_name == 'FGSM_Linf':
        steps = 1
    elif attack_name == 'Noise_Linf':
        steps = 0
    elif attack_name == 'CW_L2':
        steps = 100
    else:
        # TI / BIM / PGD / DI / MI _FGSM
        steps = 10
        pass
    eps_step = 2 * epsilon / (steps + 1e-7)

    total_labels_list = list()
    # Start to attack

    clean_images = list()
    for i, (images, target) in enumerate(tqdm(test_loader)):

        clean_images.append(images.cpu().numpy())
        if True:
            images = images.cuda()
            target = target.cuda()

        adv_images = gen_adv_samples(src_model, loss_fn, images, [target, []], 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.cpu().detach().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)
        total_labels_list.append(gt_concatnate)

    # Calculate attack ACC
    metric_counter['acc'] = accuracy_score(\
        metric_counter['gt'], \
        metric_counter['pred'], normalize=True)
    print("Acc {:.3f} Using {}".format(\
        metric_counter['acc'], attack_name))

    save_imgs = np.concatenate(metric_counter['data'], axis=0)
    base_dir = os.path.join(os.getcwd(), f'runs_{args.dataset}', args.arch)
    if not os.path.isdir(base_dir): os.mkdir(base_dir)
    np.save(f'{base_dir}/{args.attack}.npy', save_imgs)

    total_labels_list = np.concatenate(total_labels_list, axis=0)
    np.save(f'{base_dir}/gt.npy', total_labels_list)
    print(f'Save {base_dir}/{args.attack}.npy')
    clean_images = np.concatenate(clean_images, axis=0)
    np.save(f'{base_dir}/clean.npy', clean_images)
    diff = np.abs(clean_images - save_imgs) 
    diff = diff.reshape(diff.shape[0], -1)
    print(f'Difference : L_inf max {diff.max()}')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Training Code')
    parser.add_argument('--root_dir', metavar='DIR', default='/apdcephfs/share_1290796/qingsongyao/SecureMedIA/dataset/aptos2019/',
                        help='path to dataset')
    parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet50')
    parser.add_argument('-b', '--batch-size', default=32, 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('-i', '--layer_index', default=7, type=int,
                        help='Fold Number')
    parser.add_argument('--dataset', default='APTOS', type=str,
                        help='Fold Number')
    parser.add_argument('--attack', default='I_FGSM_Linf_1', type=str,
                        help='Fold Number')                        

    args = parser.parse_args()

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