} EPS = args.eps if args.eps is not None else DEFAULT_EPS[args.mode] LR = args.pgd_lr if args.pgd_lr is not None else 2*EPS/NUM_STEPS NORMALIZER = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) SAVE_ITERS = args.save_iters ATTACKS = { "l2": pgd_l2, "linf": pgd_linf, } attack = ATTACKS[MODE] loss_fn = ch.nn.CrossEntropyLoss() param_set = net.parameters() if args.opt == "sgd": opt = optim.SGD(param_set, lr=args.sgd_lr, momentum=0.9, weight_decay=2e-4) scheduler = optim.lr_scheduler.MultiStepLR(opt, milestones=[50,100,150,500], gamma=0.1) elif args.opt == "adam": opt = optim.Adam(param_set, lr=args.sgd_lr) scheduler = optim.lr_scheduler.MultiStepLR(opt, milestones=[args.num_epochs+1], gamma=0.1) # lr for mnist is 1e-4, 1e-5 elif args.opt == 'yf': opt = YFOptimizer(param_set, lr=args.sgd_lr, clip_thresh=None, adapt_clip=False) #scheduler = optim.lr_scheduler.MultiStepLR(opt, milestones=[args.num_epochs+1], gamma=0.1) def se(x1, x2, reduction='mean'): y = ch.norm((x1-x2).view(x1.shape[0],-1),dim=-1,p=2)**2 if reduction=='sum':
def encode(x, bypass=False, no_decode=False, only_decode=False): if only_decode: return ae.decode(x) if bypass: return x if not args.no_norm: x = ch.stack([NORMALIZER(x[i]) for i in range(x.shape[0])]) if args.use_orig: return x return ae(x, no_decode=no_decode) loss_fn = ch.nn.CrossEntropyLoss() if args.trainable_encode: param_set = [{'params': net.parameters()}, {'params': ae.parameters()}] elif args.trainable_decode: param_set = [{'params': net.parameters()}, {'params': ae.decode_vars}] else: param_set = net.parameters() if args.opt == "sgd": opt = optim.SGD(param_set, lr=args.sgd_lr, momentum=0.9, weight_decay=2e-4) scheduler = optim.lr_scheduler.MultiStepLR(opt, milestones=[50, 100, 150, 500], gamma=0.1) elif args.opt == "adam": opt = optim.Adam(param_set, lr=args.sgd_lr) scheduler = optim.lr_scheduler.MultiStepLR( opt, milestones=[args.num_epochs + 1], gamma=0.1) elif args.opt == 'yf':