def main(args):
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
seed = args.seed
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
np.set_printoptions(precision=4)
torch.set_printoptions(precision=4)
solver = Solver(args)
if args.mode == 'train':
solver.train()
elif args.mode == 'test':
solver.test()
elif args.mode == 'generate':
solver.generate(target=args.target,
epsilon=args.epsilon,
alpha=args.alpha,
iteration=args.iteration)
elif args.mode == 'ad_train':
solver.ad_train(target=args.target,
alpha=args.alpha,
iteration=args.iteration,
lamb=0.3)
elif args.mode == 'ad_test': #adversarial image test
solver.ad_test(target=args.target,
epsilon=args.epsilon,
alpha=args.alpha,
iteration=args.iteration)
else:
return
print('[*] Finished')
