复现 attack as defense

attack_as_defense_detector.py

import argparse
from get_data import get_dataset
from train_model import MNIST, train_model_iteration
import random
import numpy as np
import torch

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)  # cpu
    torch.cuda.manual_seed(seed)  # gpu
    torch.cuda.manual_seed_all(seed)  # all gpus

def main(arg):
    dataset_name= arg.dataset
    detector_method = arg.detector
    attack_method = arg.attack
    #加载数据
    train_data, test_data = get_dataset(dataset_name, 32)
    #加载模型
    def try_gpu(i=0):
        if torch.cuda.device_count() >= i + 1:
            print('gpu is avaliable',f'cuda:{i}')
            return torch.device(f'cuda:{i}')
        else:
            return torch.device('cpu')
    mnist_model = MNIST()
    train_model_iteration(mnist_model, train_data, test_data, 10, 0.01, try_gpu())
    #生成对抗样本

    #生成对抗代价
    pass

if __name__ == '__main__':
    parse = argparse.ArgumentParser()
    parse.add_argument('--dataset', help='Dataset to use; either minist, cifar', required=False, type=str, default='mnist')
    parse.add_argument('--init', help='If this is the first time to run this script, need to generate the model and training data', action='store_true')
    parse.add_argument('-d', '--detector', help='Detector type; either knn or zscore', required=False, type=str, default='knn')
    parse.add_argument('-a', '--attack', help='Attack to use; recommand to sue JSMA, BIM', required=False, type=str, default='JSMA')

    args = parse.parse_args()
    main(args)

train_model.py

import torch
import tqdm
from torch import nn
from torch.nn import functional as F

class MNIST(nn.Module):
    def __init__(self):
        super(MNIST, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=(5,5))
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(5,5))
        self.flatten = nn.Flatten()
        self.linear1 = nn.Linear(4*4*32, 512)
        self.linear2 = nn.Linear(512, 10)

    def forward(self, x):
        out = x.view(-1, 1,28,28)
        out = self.conv1(out)
        out = F.relu(out)
        out = F.max_pool2d(out, [2,2])
        out = self.conv2(out)
        out = F.relu(out)
        out = F.max_pool2d(out, [2,2])

        out = self.flatten(out)
        out = self.linear1(out)
        out = F.relu(out)
        out = self.linear2(out)
        return out


def train_model_iteration(model:nn.Module, train_data:torch.Tensor, valid_data:torch.Tensor, epochs, lr, device):
    def evaluate_accuracy_gpu(model, data, device=None):
        if isinstance(model, torch.nn.Module):
            model.eval()
            if not device:
                device = next(iter(model.parameters())).device
        result = []
        for value, label in data:
            x = value.to(device)
            y = label.to(device)
            y_pre = model(x)
            result.append(y==y_pre)
        return torch.sum(torch.tensor(result))/len(data)
    def adjust_learning_rate(optimizer, epoch):
        """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
        if epoch < 10:
            lr = 0.01
        elif epoch > 10 and epoch < 50:
            lr = 0.001
        else:
            lr = 1e-5
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

    def init_weight(m):
        if type(m) == nn.Linear or type(m) == nn.Conv2d:
            nn.init.xavier_uniform_(m.weight)
    model.apply(init_weight)
    print('trainign on ', device)
    model.to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    loss = nn.CrossEntropyLoss()
    for epoch in range(epochs):
        sum_loss = 0.0
        train_correct = 0.0
        adjust_learning_rate(optimizer, epoch)
        model.train()
        pbar = tqdm.tqdm(train_data)
        for i, batch in enumerate(pbar):
            optimizer.zero_grad()
            inputs, labels = batch
            X = inputs.to(device)
            y = labels.to(device)
            outputs = model(X)
            l = loss(outputs, y)
            l.backward()
            pbar.set_description(f'Loss: {l.item():.04f}')
            optimizer.step()
            _, index = torch.max(outputs.data, dim=1)
            sum_loss += l.data  # 累加loss
            train_correct += torch.sum(torch.tensor(id == y.data))
        valid_acc = evaluate_accuracy_gpu(model, valid_data, device)
        print('avrage_loss:%.03f; train_acc:%0.3f; valid_acc:%0.3f'
              % (sum_loss / len(train_data), 100 * train_correct / len(train_data),valid_acc))

get_data.py

import pathlib
import torch

def get_dataset(dataset, batchsize=32):
    from torchvision import datasets, transforms

    def transform(i):
        data = torch.tensor((i.getdata()))
        data = data / 255.0
        return data

    if dataset == "cifar":
        if not pathlib.Path('./data/CIFAR').exists():
            download = True
            print('第一次加载数据，需要远程拉取，注意联网！！')
        else:
            download = False
        data_train = datasets.CIFAR10(root='./data/CIFAR', train=True, transform=transform, download= download)
        data_test = datasets.CIFAR10(root='./data/CIFAR', train=False, transform=transform, download= download)
    elif dataset == 'mnist':
        if not pathlib.Path('./data/MNIST').exists():
            download = True
            print('第一次加载数据，需要远程拉取，注意联网！！')
        else:
            download = False
        data_train = datasets.MNIST(root='./data/MNIST', train=True, transform=transform, download= download)
        data_test = datasets.MNIST(root='./data/MNIST', train=False, transform=transform, download= download)
    else:
        raise Exception("Dataset not implemented")

    from torch.utils.data import DataLoader
    data_train_loader = DataLoader(dataset=data_train, batch_size=batchsize, shuffle=True)
    data_test_loader = DataLoader(dataset=data_test, batch_size=batchsize, shuffle=True)
    return data_train_loader, data_test_loader