feaSca_main.py

import os,argparse,sys,datetime,time

import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, TensorDataset

from models import *
from attack_methods import *
from ot import *
from training_config import *
from utils import *
from use_attacks import *
from attacks import *

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


if __name__ == '__main__':
    
    parser = argparse.ArgumentParser(description='MNIST')

    # model hyper-parameter variables
    parser.add_argument('--itr', default=100, metavar='itr', type=int, help='no. of iteration for training')
    parser.add_argument('--eps', default=0.15, metavar='eps', type=float, help='epsilon value')
    parser.add_argument('--max_iter', default=10, metavar='max_iter', type=int, help='maximum iteration for attack convergence')
    parser.add_argument('--step_size', default=0.08, metavar='step_size', type=float, help='step size in each iteration of attack')
    
    args = parser.parse_args()


num_epochs = args.itr
eps = args.eps
max_iter = args.max_iter
step_size = args.step_size


trainset = datasets.MNIST(root='/home/aminul/data', train = True, download =False, transform = transforms.ToTensor())
testset = datasets.MNIST(root='/home/aminul/data', train = False, download =False, transform = transforms.ToTensor())
num_classes = 10


train_loader = DataLoader(trainset, 100, shuffle =True, num_workers = 2)
test_loader = DataLoader(testset, 100, shuffle = False, num_workers = 2)


net = WideResNet(depth=28,
                           num_classes=num_classes,
                           widen_factor=10)
net = net.to(device)
net = nn.DataParallel(net, device_ids=[0, 1, 2])


criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(),lr=0.1)


for epoch in range(0, num_epochs):
    FeaScatter_train(net,train_loader,optimizer,criterion,epoch,eps,1,step_size,device) # 1 is the max_iter
    best_acc = test(net,test_loader,optimizer,criterion,epoch,device,"fea_sca_adversarial_training_")


print("\n\n\n----------------")
print("Loading Best Model")
checkpoint = torch.load('./checkpoint/fea_sca_adversarial_training_model.t7')
net = checkpoint['net']
print("Done !")

FGSM_test(net,test_loader,optimizer,criterion,eps,device)    
PGD_test(net,test_loader,optimizer,criterion,eps,max_iter,step_size,device)