with open(os.path.join(args.load, args.attack+"_"+str(args.eps), "adversary.pkl"), "wb") as f :
pickle.dump(adv_saver,f)
logger.info("Attacked Accuracy : {:.4f}".format(adv_acc))
logger.info("Attacked Loss : {:.4f}".format(adv_loss))
logger.info("Finished")
logger.info("="*80)
if __name__ == "__main__" :
args.device = torch.device("cuda:" + str(args.gpu) if torch.cuda.is_available() else "cpu")
logger = init_logger(logpath=args.load, experiment_name="attack-"+str(args.attack)+"-"+str(args.eps))
in_channels = 1 if args.eval=="mnist" else 3
if args.eval == "mnist" or args.eval == "norm" :
from model.mnist import mnist_model
model = mnist_model(args.model, layers=args.block, norm_type=args.norm)
elif args.eval == "cifar10" :
from model.cifar10 import cifar_model
model = cifar_model(args.model, layers=args.block, norm_type=args.norm)
logger.info(args)
logger.info(model)
model.to(args.device)
if args.crit == "acc" :
model_dict = torch.load(os.path.join(args.load,"model_acc.pt"), map_location=str(args.device))
elif args.crit == "last" :
model_dict = torch.load(os.path.join(args.load,"model_final.pt"), map_location=str(args.device))
elif args.crit == "loss" :
model_dict = torch.load(os.path.join(args.load, "model_loss.pt"), map_location=str(args.device))
model.load_state_dict(model_dict["state_dict"], strict=False)
model.to(args.device)
if __name__ == "__main__" :
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
np.random.seed(args.seed)
if os.path.exists(args.save) :
raise NameError("previous experiment '{}' already exists!".format(args.save))
os.makedirs(args.save)
logger = init_logger(logpath=args.save, experiment_name="logs-"+args.model)
logger.info(args)
if args.gpu >= 0 :
args.device = torch.device("cuda:" + str(args.gpu) if torch.cuda.is_available() else "cpu")
else :
args.device = torch.device("cpu")
model = mnist_model(args.model, layers=args.block, norm_type=args.norm, init_option=args.init)
logger.info(model)
model.to(args.device)
train_loader, test_loader, train_eval_loader = get_mnist_loaders()
loader = {"train_loader": train_loader, "train_eval_loader": train_eval_loader, "test_loader": test_loader}
if args.opt == "sgd" :
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[60,100,140], gamma=0.1)
elif args.opt == "adam" :
optimizer = torch.optim.Adam(model.parameters())
scheduler = None
elif args.opt == "rms" :
optimizer = torch.optim.RMSprop(model.parameters(), lr=1e-3)
scheduler = None