-
Notifications
You must be signed in to change notification settings - Fork 2
/
easy_train.py
executable file
·103 lines (88 loc) · 3.02 KB
/
easy_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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
#!/usr/bin/env python
import os
import argparse
import torch as ch
from torchvision import transforms
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from robustness import model_utils, datasets, train, defaults
from cox.utils import Parameters
import cox.store
from data_utils import *
parser = argparse.ArgumentParser(description='train configs')
parser.add_argument('--dataset' ,help='the used dataset')
parser.add_argument('--arch' , default='resnet18', help='the used model')
parser.add_argument('--freq', default=None, help='frequecy')
parser.add_argument('--r', default=0, help='radius of frequecy domain')
parser.add_argument('--lr', default=0.1, help='leanring rate')
parser.add_argument('--epoch', default=200, help='epoch')
parser.add_argument('--task', default=None, help='expriment id')
parser.add_argument('--gpu', default="0", help='choose gpu to train')
args = parser.parse_args()
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
DATA_PATH_DICT = {'cifar10': '/data/cifar','cifar100':'/data/cifar','resIN':'/data/ILSVRC2012'}
DATA_NAME_DICT = {'cifar10':'CIFAR', 'cifar100':'CIFAR100','resIN':'RestrictedImageNet'}
DEVICES_ID = None
if len(args.gpu)>1:
DEVICES_ID = [int(id) for id in args.gpu.split(',')]
# configs
ARCH = args.arch
FREQ = args.freq
RADIUS = args.r
DATA = args.dataset
OUT_DIR = args.dataset
LR = float(args.lr)
EPOCH = int(args.epoch)
NUM = int(args.number)
dataset_function = getattr(datasets, DATA_NAME_DICT[DATA])
dataset = dataset_function(DATA_PATH_DICT[DATA])
# DATA loader
print("preparing data {}...".format(DATA))
train_loader, test_loader = get_dataloader(DATA, FREQ, RADIUS)
# train_loader, test_loader = dataset.make_loaders(workers=8, batch_size=128)
print('data loader is ok :)')
# prepare model
print("preparing model {}...".format(ARCH))
model_kwargs = {
'arch': ARCH,
'dataset': dataset
}
model, _ = model_utils.make_and_restore_model(**model_kwargs)
print('model is ok :)')
#train
train_kwargs = {
'out_dir' : OUT_DIR,
'epochs' : EPOCH,
'lr' : LR,
'momentum' : 0.9,
'weight_decay': 5e-4,
'lr_interpolation ': 'step',
'custom_lr_multplier ': '[(60,0.1),(120,0.1),(180,0.1)]',
'save_ckpt_iters' : 1,
'adv_train': 0,
'constraint': '2',
'eps': 0.5,
'attack_lr': 0.1,
'attack_steps':7,
'use_best': True,
'random_restarts': 0
}
train_args = Parameters(train_kwargs)
train_args = defaults.check_and_fill_args(train_args,
defaults.TRAINING_ARGS, dataset_function)
#log
if args.task==None:
if FREQ==None:
task = ARCH+'_baseline'
else:
task = ARCH+'_'+FREQ+'_'+RADIUS
else:
task = args.task
out_store = cox.store.Store(OUT_DIR, exp_id = task)
print("="*20)
if FREQ:
print("Ready to train {} on dataset {} ,frequency : {} and r: {}".format(ARCH,DATA,FREQ,RADIUS))
else:
print("Ready to train {} on dataset {} , baseline".format(ARCH,DATA))
print("="*20)
train.train_model(train_args, model, (train_loader, test_loader), store=out_store, dp_device_ids=DEVICES_ID)