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) and args.load == "none" :
raise NameError("previous experiment '{}' already exists!".format(args.save))
if args.load == "none" :
os.makedirs(args.save)
logger = init_logger(logpath=args.save, experiment_name="logs-"+args.model)
logger.info(args)
args.device = torch.device("cuda:" + str(args.gpu) if torch.cuda.is_available() else "cpu")
train_loader, test_loader, train_eval_loader = get_cifar10_loaders(data_aug=True, batch_size=args.tbsize)
model = cifar_model(args.model, layers=args.block, norm_type=args.norm, init_option=args.init)
logger.info(model)
if args.load != "none" :
model.load_state_dict(torch.load(os.path.join(args.load, "model_final.pt"), map_location=args.device)['state_dict'])
model.to(args.device)
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, weight_decay=args.decay, momentum=0.9, nesterov=args.nesterov)
if args.adv == "none" :
if args.model == "ssp2" or args.model == "ssp3" :
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[70,120,160], gamma=0.1)
else :
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[60,100,140], gamma=0.1)
if args.epochs <= 100 :
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30,60,90], gamma=0.1)
args.save = os.path.join("experiments", args.save)
if os.path.exists(args.save) and args.load == "none":
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)
args.device = torch.device(
"cuda:" + str(args.gpu) if torch.cuda.is_available() else "cpu")
train_loader, test_loader, train_eval_loader = get_cifar10_loaders(
data_aug=True, batch_size=args.tbsize)
model = cifar_model(args.model, layers=args.block, norm_type=args.norm)
logger.info(model)
model.to(args.device)
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,
weight_decay=args.decay,
momentum=0.9,
nesterov=args.nesterov)
if args.adv == "none":