make_patch_multi.py

import argparse
import os
import random
import numpy as np
import warnings

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torch.nn.functional as F
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from torch.utils.data.sampler import SubsetRandomSampler

from pretrained_models_pytorch import pretrainedmodels

from utils import *

parser = argparse.ArgumentParser()
parser.add_argument('--workers', type=int,
                    help='number of data loading workers', default=2)
parser.add_argument('--epochs', type=int, default=20,
                    help='number of epochs to train for')
parser.add_argument('--cuda', action='store_true', help='enables cuda')

parser.add_argument('--target', type=int, default=859,
                    help='The target class: 859 == toaster')
parser.add_argument('--conf_target', type=float, default=0.9,
                    help='Stop attack on image when target classifier reaches this value for target class')

parser.add_argument('--max_count', type=int, default=1000,
                    help='max number of iterations to find adversarial example')
parser.add_argument('--patch_type', type=str, default='circle',
                    help='patch type: circle or square')
parser.add_argument('--patch_size', type=float, default=0.05,
                    help='patch size. E.g. 0.05 ~= 5% of image ')

parser.add_argument('--train_size', type=int, default=2000,
                    help='Number of training images')
parser.add_argument('--test_size', type=int, default=2000,
                    help='Number of test images')

parser.add_argument('--image_size', type=int, default=299,
                    help='the height / width of the input image to network')

parser.add_argument('--plot_all', type=int, default=1,
                    help='1 == plot all successful adversarial images')

parser.add_argument('--netClassifier', default='inceptionv3',
                    help="The target classifier")

parser.add_argument('--outf', default='./logs',
                    help='folder to output images and model checkpoints')
parser.add_argument('--manualSeed', type=int, default=1338, help='manual seed')

opt = parser.parse_args()
print(opt)
warnings.simplefilter("ignore")

try:
    os.makedirs(opt.outf)
except OSError:
    pass

try:
    os.makedirs(f'{opt.outf}/train')
except OSError:
    pass

try:
    os.makedirs(f'{opt.outf}/test')
except OSError:
    pass

if opt.manualSeed is None:
    opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
np.random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
device = torch.device('mps')

# if opt.cuda:
#    torch.cuda.manual_seed_all(opt.manualSeed)

cudnn.benchmark = True

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

target = opt.target
conf_target = opt.conf_target
max_count = opt.max_count
patch_type = opt.patch_type
patch_size = opt.patch_size
image_size = opt.image_size
train_size = opt.train_size
test_size = opt.test_size
plot_all = opt.plot_all

assert train_size + \
    test_size <= 50000, "Traing set size + Test set size > Total dataset size"

print("=> creating model ")
netClassifier = pretrainedmodels.__dict__[opt.netClassifier](
    num_classes=1000, pretrained='imagenet')

netClassifiers_names = ['inceptionv3', 'resnet50', 'vgg16', 'vgg19']

netClassifiers = [pretrainedmodels.__dict__[name](
    num_classes=1000, pretrained='imagenet') for name in netClassifiers_names]

if opt.cuda:
    netClassifier.to(device)
    # netClassifier.cuda()


print('==> Preparing data..')
normalize = transforms.Normalize(mean=netClassifier.mean,
                                 std=netClassifier.std)
idx = np.arange(50000)
np.random.shuffle(idx)
training_idx = idx[:train_size]
test_idx = idx[train_size:(test_size+train_size)]

train_loader = torch.utils.data.DataLoader(
    dset.ImageFolder('./imagenetdata/val', transforms.Compose([
        transforms.Resize(round(max(netClassifier.input_size)*1.050)),
        transforms.CenterCrop(max(netClassifier.input_size)),
        transforms.ToTensor(),
        ToSpaceBGR(netClassifier.input_space == 'BGR'),
        ToRange255(max(netClassifier.input_range) == 255),
        normalize,
    ])),
    batch_size=1, shuffle=False, sampler=SubsetRandomSampler(training_idx),
    num_workers=opt.workers, pin_memory=True)

test_loader = torch.utils.data.DataLoader(
    dset.ImageFolder('./imagenetdata/val', transforms.Compose([
        transforms.Resize(round(max(netClassifier.input_size)*1.050)),
        transforms.CenterCrop(max(netClassifier.input_size)),
        transforms.ToTensor(),
        ToSpaceBGR(netClassifier.input_space == 'BGR'),
        ToRange255(max(netClassifier.input_range) == 255),
        normalize,
    ])),
    batch_size=1, shuffle=False, sampler=SubsetRandomSampler(test_idx),
    num_workers=opt.workers, pin_memory=True)

min_in, max_in = netClassifier.input_range[0], netClassifier.input_range[1]
min_in, max_in = np.array([min_in, min_in, min_in]
                          ), np.array([max_in, max_in, max_in])
mean, std = np.array(netClassifier.mean), np.array(netClassifier.std)
min_out, max_out = np.min((min_in-mean)/std), np.max((max_in-mean)/std)


def train(epoch, patch, patch_shape):
    for i in range(len(netClassifiers)):
        netClassifiers[i].eval()
    success = 0
    total = 0
    recover_time = 0
    try:
        os.makedirs(f'{opt.outf}/train/{epoch}')
    except OSError:
        pass
    for batch_idx, (data, labels) in enumerate(train_loader):
        for i, classifier in enumerate(netClassifiers):
            if opt.cuda:
                data = data.cuda()
                labels = labels.cuda()
            data, labels = Variable(data), Variable(labels)

            prediction = classifier(data)

            # only computer adversarial examples on examples that are originally classified correctly
            if prediction.data.max(1)[1][0] != labels.data[0]:
                print(
                    f'{netClassifiers_names[i]} predicted incorectly {prediction.data.max(1)[1][0]} != {labels.data[0]}')
                continue

            total += 1

            # transform path
            data_shape = data.data.cpu().numpy().shape
            if patch_type == 'circle':
                patch, mask, patch_shape = circle_transform(
                    patch, data_shape, patch_shape, image_size)
            elif patch_type == 'square':
                patch, mask = square_transform(
                    patch, data_shape, patch_shape, image_size)
            patch, mask = torch.FloatTensor(patch), torch.FloatTensor(mask)
            if opt.cuda:
                patch, mask = patch.to(device), mask.to(device)
                #patch, mask = patch.cuda(), mask.cuda()
            patch, mask = Variable(patch), Variable(mask)

            adv_x, mask, patch = attack(data, patch, mask, classifier, i)

            adv_label = classifier(adv_x).data.max(1)[1][0]
            ori_label = labels.data[0]

            if adv_label == target:
                success += 1

                if plot_all == 1:
                    # plot source image
                    vutils.save_image(data.data, "./%s/train/%d/%d_%d_original_%s.png" %
                                      (opt.outf, epoch, batch_idx, ori_label, netClassifiers_names[i]), normalize=True)

                    # plot adversarial image
                    vutils.save_image(adv_x.data, "./%s/train/%d/%d_%d_adversarial_%s.png" %
                                      (opt.outf, epoch, batch_idx, adv_label, netClassifiers_names[i]), normalize=True)

            masked_patch = torch.mul(mask, patch)
            patch = masked_patch.data.cpu().numpy()
            new_patch = np.zeros(patch_shape)
            for i in range(new_patch.shape[0]):
                for j in range(new_patch.shape[1]):
                    new_patch[i][j] = submatrix(patch[i][j])

            patch = new_patch

        # log to file
        keeper = False
        if total == 0:
            keeper = True
            total = 1
        progress_bar(batch_idx, len(train_loader),
                     "Train Patch Success: {:.3f}".format(success/total))
        if keeper:
            total = 0

    return patch


def test(epoch, patch, patch_shape):
    for i in range(len(netClassifiers)):
        netClassifiers[i].eval()
    success = 0
    total = 0
    try:
        os.makedirs(f'{opt.outf}/test/{epoch}')
    except OSError:
        pass
    for batch_idx, (data, labels) in enumerate(test_loader):
        for i, classifier in enumerate(netClassifiers):
            if opt.cuda:
                data = data.to(device)
                labels = labels.to(device)
                #data = data.cuda()
                #labels = labels.cuda()
            data, labels = Variable(data), Variable(labels)

            prediction = classifier(data)

            # only computer adversarial examples on examples that are originally classified correctly
            if prediction.data.max(1)[1][0] != labels.data[0]:
                continue

            total += 1

            # transform path
            data_shape = data.data.cpu().numpy().shape
            if patch_type == 'circle':
                patch, mask, patch_shape = circle_transform(
                    patch, data_shape, patch_shape, image_size)
            elif patch_type == 'square':
                patch, mask = square_transform(
                    patch, data_shape, patch_shape, image_size)
            patch, mask = torch.FloatTensor(patch), torch.FloatTensor(mask)
            if opt.cuda:
                patch, mask = patch.to(device), mask.to(device)
                #patch, mask = patch.cuda(), mask.cuda()
            patch, mask = Variable(patch), Variable(mask)

            adv_x = torch.mul((1-mask), data) + torch.mul(mask, patch)
            adv_x = torch.clamp(adv_x, min_out, max_out)

            adv_label = classifier(adv_x).data.max(1)[1][0]
            ori_label = labels.data[0]

            if adv_label == target:
                success += 1

            vutils.save_image(data.data, "./%s/test/%d/%d_%d_original_%s.png" %
                              (opt.outf, epoch, batch_idx, ori_label, netClassifiers_names[i]), normalize=True)
            vutils.save_image(adv_x.data, "./%s/test/%d/%d_%d_adversarial_%s.png" %
                              (opt.outf, epoch, batch_idx, adv_label, netClassifiers_names[i]), normalize=True)

            masked_patch = torch.mul(mask, patch)
            patch = masked_patch.data.cpu().numpy()
            new_patch = np.zeros(patch_shape)
            for i in range(new_patch.shape[0]):
                for j in range(new_patch.shape[1]):
                    new_patch[i][j] = submatrix(patch[i][j])

            patch = new_patch

        # log to file
        progress_bar(batch_idx, len(test_loader),
                     "Test Success: {:.3f}".format(success/total))


def attack(x, patch, mask, classifier, i):
    classifier.eval()

    x_out = F.softmax(classifier(x))
    target_prob = x_out.data[0][target]

    adv_x = torch.mul((1-mask), x) + torch.mul(mask, patch)

    count = 0

    while conf_target > target_prob:
        count += 1
        adv_x = Variable(adv_x.data, requires_grad=True)
        adv_out = F.log_softmax(classifier(adv_x))

        adv_out_probs, adv_out_labels = adv_out.max(1)

        Loss = -adv_out[0][target]
        Loss.backward()

        adv_grad = adv_x.grad.clone()

        adv_x.grad.data.zero_()

        patch -= adv_grad

        adv_x = torch.mul((1-mask), x) + torch.mul(mask, patch)
        adv_x = torch.clamp(adv_x, min_out, max_out)

        out = F.softmax(classifier(adv_x))
        target_prob = out.data[0][target]
        #y_argmax_prob = out.data.max(1)[0][0]

        #print(count, conf_target, target_prob, y_argmax_prob)

        if count >= opt.max_count:
            print(
                f'{netClassifiers_names[i]} over max count')
            break

    return adv_x, mask, patch


if __name__ == '__main__':
    if patch_type == 'circle':
        patch, patch_shape = init_patch_circle(image_size, patch_size)
    elif patch_type == 'square':
        patch, patch_shape = init_patch_square(image_size, patch_size)
    else:
        sys.exit("Please choose a square or circle patch")

    for epoch in range(1, opt.epochs + 1):
        patch = train(epoch, patch, patch_shape)
        test(epoch, patch, patch_shape)