-
Notifications
You must be signed in to change notification settings - Fork 4
/
config.py
64 lines (49 loc) · 2.3 KB
/
config.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 argparse
def str2bool(v):
return v.lower() in ('true', '1')
def add_argument_group(name):
arg = parser.add_argument_group(name)
arg_lists.append(arg)
return arg
arg_lists = []
parser = argparse.ArgumentParser()
# Network
net_arg = add_argument_group('Network')
net_arg.add_argument('--model', type=str, default='vit', help='Model name')
# Optimizer arguments
opt_arg = add_argument_group('Optimizer')
opt_arg.add_argument('--lr', type=float, default=1e-2)
# Scheduler
opt_arg.add_argument('--max_epochs', type=int, default=50)
# Directories
dir_arg = add_argument_group('Directories')
dir_arg.add_argument('--log_dir', type=str, default='outputs/default')
# Data
data_arg = add_argument_group('Data')
data_arg.add_argument('--dataset', type=str, default='Synapse')
data_arg.add_argument('--batch_size', type=int, default=2)
data_arg.add_argument('--val_batch_size', type=int, default=2)
data_arg.add_argument('--test_batch_size', type=int, default=2)
data_arg.add_argument('--num_workers', type=int, default=4)
data_arg.add_argument('--num_val_workers', type=int, default=4)
data_arg.add_argument('--ignore_index', type=int, default=255)
data_arg.add_argument('--cityscapes_path',type=str,default='/images/PublicDataset/cityscapes')
data_arg.add_argument('--synapse_path',type=str,default='./data/Synapse')
# Training / test parameters
train_arg = add_argument_group('Training')
train_arg.add_argument('--is_train', type=str2bool, default=True)
train_arg.add_argument('--stat_freq', type=int, default=40, help='print frequency')
train_arg.add_argument('--save_freq', type=int, default=1000, help='save frequency')
train_arg.add_argument('--val_freq', type=int, default=1000, help='validation frequency')
train_arg.add_argument('--train_phase', type=str, default='train', help='Dataset for training')
train_arg.add_argument('--val_phase', type=str, default='validation', help='Dataset for validation')
train_arg.add_argument('--resume', default=None, type=str, help='path to latest checkpoint')
# Test
test_arg = add_argument_group('Test')
test_arg.add_argument('--test_phase', type=str, default='test', help='Dataset for test')
# Misc
misc_arg = add_argument_group('Misc')
misc_arg.add_argument('--is_cuda', type=str2bool, default=True)
def get_config():
config = parser.parse_args()
return config # Training settings