-
Notifications
You must be signed in to change notification settings - Fork 2
/
mnist_train.py
64 lines (55 loc) · 2.66 KB
/
mnist_train.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
import os
import argparse
import time
import torch
import torch.nn as nn
import numpy as np
from dataloader import get_mnist_loaders, get_cifar10_loaders, inf_generator
from utils import init_logger, RunningAverageMeter, accuracy
from model.mnist import mnist_model
from container import trainer
parser = argparse.ArgumentParser("MNIST")
parser.add_argument("--model", type=str, default="res", choices=["res", "ssp2", "ssp3", "ark"])
parser.add_argument("--epochs", type=int, default=200)
parser.add_argument("--lr", type=float, default=0.1)
parser.add_argument("--block", type=eval, default=5)
parser.add_argument("--hist", type=eval, default=False)
parser.add_argument("--norm", type=str, default="g")
parser.add_argument("--save", type=str, default="exp")
parser.add_argument("--gpu", type=int, default=0)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--adv", type=str, default="none", choices=["none", "fgsm", "pgd", "ball"])
parser.add_argument("--iters", type=int, default=10)
parser.add_argument("--init", type=str, default="basic")
parser.add_argument("--opt", type=str, default="sgd", choices=["sgd", "adam", "rms"])
parser.add_argument("--repeat", type=int, default=5)
args = parser.parse_args()
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
adv_train = args.adv if args.adv != "none" else None
model = trainer(model, logger, loader, args, "mnist", optimizer, scheduler, adv_train=adv_train)