DeepCell/config.py at main · cure-lab/DeepCell · GitHub

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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
import os
import argparse

def get_parse_args():
    parser = argparse.ArgumentParser(description='Pytorch training script of DeepGate.')

    # basic experiment setting
    parser.add_argument('--exp_id', default='train')
    parser.add_argument('--debug', action='store_true', default=False)
    parser.add_argument('--load_model', default='',
                             help='path to pretrained model')
    parser.add_argument('--resume', action='store_true',
                             help='resume an experiment. '
                                  'Reloaded the optimizer parameter and '
                                  'set load_model to model_last.pth '
                                  'in the exp dir if load_model is empty.')

    # experiment
    parser.add_argument('--disable_encode', action='store_true', default=False)
    parser.add_argument('--refine', default='aig', choices=['aig', 'pm'])
    parser.add_argument('--wo_view', action='store_true', default=False)

    # GNN
    parser.add_argument('--pm_aggr', default='dg2', choices=['dg2', 'gat', 'gcn'],)
    parser.add_argument('--aig_encoder', default='dg2', choices=['dg2', 'pg', 'dg3', 'gcn', 'hoga'])

    # system
    parser.add_argument('--gpus', default='-1',
                             help='-1 for CPU, use comma for multiple gpus')
    parser.add_argument("--local-rank", default=0, type=int)
    parser.add_argument('--num_workers', type=int, default=4,
                             help='dataloader threads. 0 for single-thread.')
    parser.add_argument('--not_cuda_benchmark', action='store_true',
                             help='disable when the input size is not fixed.')
    parser.add_argument('--random-seed', type=int, default=208,
                             help='random seed')

    # log
    parser.add_argument('--print_iter', type=int, default=0,
                             help='disable progress bar and print to screen.')
    parser.add_argument('--hide_data_time', action='store_true',
                             help='not display time during training.')
    parser.add_argument('--save_all', action='store_true',
                             help='save model to disk every 5 epochs.')
    parser.add_argument('--save_intervals', type=int, default=5,
                             help='number of epochs to run validation.')
    parser.add_argument('--metric', default='loss',
                             help='main metric to save best model')

    # dataset settings
    parser.add_argument('--data_dir', default='./data/train',
                             type=str, help='the path to the dataset')
    parser.add_argument('--max_token_size', default=4096, type=int,
                             help='max token size for each circuit pair')

    # train and val
    parser.add_argument('--lr', type=float, default=1.0e-4,
                             help='learning rate for batch size 32.')
    parser.add_argument('--weight_decay', type=float, default=1e-10,
                             help='weight decay (default: 1e-10)')
    parser.add_argument('--lr_step', type=str, default='40',
                             help='drop learning rate by 10.')
    parser.add_argument('--num_epochs', type=int, default=40,
                             help='total training epochs.')
    parser.add_argument('--batch_size', type=int, default=4,
                             help='batch size')
    parser.add_argument('--trainval_split', default=0.9, type=float,
                             help='the splitting setting for training dataset and validation dataset.')

    # Model
    parser.add_argument('--dim_hidden', type=int, default=128)
    parser.add_argument('--tf_head', type=int, default=8)
    parser.add_argument('--tf_layer', type=int, default=4)
    parser.add_argument('--mask_ratio', type=float, default=0.15)
    parser.add_argument('--k_hop', type=int, default=4)
    parser.add_argument('--linformer', action='store_true', default=False,
                             help='use linformer instead of transformer')

    args = parser.parse_args()

    args.gpus_str = args.gpus
    args.gpus = [int(gpu) for gpu in args.gpus.split(',')]
    args.lr_step = [int(i) for i in args.lr_step.split(',')]

    if len(args.gpus) == 1:
        args.distributed = False
        if args.gpus[0] == -1:
            args.device = 'cpu'
        else:
            args.device = 'cuda:{}'.format(args.gpus[0])
    else:
        args.distributed = True
        args.device = 'distributed'

    if args.wo_view:
        args.mask_ratio = 0

    # update data settings
    args.gate_to_index = {'PI': 0, 'AND': 1, 'NOT': 2}
    args.num_gate_types = len(args.gate_to_index)
    args.dim_node_feature = len(args.gate_to_index)

    if args.debug > 0:
        args.num_workers = 0
        # args.batch_size = 1
        args.gpus = [args.gpus[0]]

    # dir
    args.root_dir = os.path.join(os.path.dirname(__file__), '.')
    args.exp_dir = os.path.join(args.root_dir, 'exp')
    args.save_dir = os.path.join(args.exp_dir, args.exp_id)
    args.debug_dir = os.path.join(args.save_dir, 'debug')
    print('The output will be saved to ', args.save_dir)

    if args.resume and args.load_model == '':
        model_path = args.save_dir
        args.load_model = os.path.join(model_path, 'model_last.pth')
    elif args.load_model != '':
        model_path = args.save_dir
        args.load_model = os.path.join(model_path, args.load_model)

    args.local_rank = 0


    return args

def update_dir(args, exp_id):
    # dir
    args.root_dir = os.path.join(os.path.dirname(__file__), '..')
    args.exp_dir = os.path.join(args.root_dir, 'exp', args.task)
    args.save_dir = os.path.join(args.exp_dir, args.exp_id)
    args.debug_dir = os.path.join(args.save_dir, 'debug')
    print('The output will be saved to ', args.save_dir)

    if args.resume and args.load_model == '':
        model_path = args.save_dir
        args.load_model = os.path.join(model_path, 'model_last.pth')
    elif args.load_model != '':
        model_path = args.save_dir
        args.load_model = os.path.join(model_path, args.load_model)

    return args