-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathconfig.py
249 lines (212 loc) · 12.8 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
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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
import os
import argparse
def get_parse_args():
parser = argparse.ArgumentParser(description='Pytorch training script of DeepGate.')
# basic experiment setting
parser.add_argument('task', default='deepsat', choices=['deepsat', 'neurosat', 'circuitsat'],
help='deepsat | neurosat | circuitsat')
parser.add_argument('--exp_id', default='default')
parser.add_argument('--debug', type=int, default=0,
help='level of visualization.'
'1: only show the final results compared with C1'
'2: debug the network gradients')
# '3: use matplot to display' # useful when lunching training with ipython notebook
# '4: save all visualizations to disk')
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.')
# system
parser.add_argument('--gpus', default='-1',
help='-1 for CPU, use comma for multiple gpus')
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('--test_data_dir', default=None,
type=str, help='the path to the testing dataset')
# circuit
parser.add_argument('--gate_types', default='*', type=str,
metavar='LIST', help='gate types in the circuits. For aig: INPUT,AND,NOT, For Circuit-sat: INPUT,AND,OR,NOT')
parser.add_argument('--dim_edge_feature', default=16,
type=int, help='the dimension of node features')
parser.add_argument('--small_train', default=0,
type=int,help='if True, use a smaller version of train set')
parser.add_argument('--un_directed', default=False, action='store_true',
help='If true, model the circuit as the undirected graph. Default: circuit as DAG')
# sat
parser.add_argument('--n_pairs', default=10000, type=int,
help='number of sat/unsat problems to generate')
parser.add_argument('--min_n', type=int, default=3,
help='min number of variables used for training')
parser.add_argument('--max_n', type=int, default=10,
help='max number of variables used for training')
# neurosat
parser.add_argument('--p_k_2', type=float, default=0.3)
parser.add_argument('--p_geo', type=float, default=0.4)
# circuitsat
parser.add_argument('--exp_depth', type=int, default=3)
# deepsat
parser.add_argument('--use_aig', action='store_true',
help='whether to use AIG.')
parser.add_argument('--no_aigopt', action='store_true',
help='whether to disable AIG optimization.')
# model settings
parser.add_argument('--activation_layer', default='relu', type=str, choices=['relu', 'relu6', 'sigmoid'],
help='The activation function to use in the FC layers.')
parser.add_argument('--norm_layer', default='batchnorm', type=str,
help='The normalization function to use in the FC layers.')
parser.add_argument('--num_fc', default=3, type=int,
help='The number of FC layers')
# recgnn / neurosat
parser.add_argument('--num_aggr', default=3, type=int,
help='the number of aggregation layers.')
parser.add_argument('--aggr_function', default='deepset', type=str, choices=['deepset', 'agnnconv', 'gated_sum', 'conv_sum'],
help='the aggregation function to use.')
parser.add_argument('--update_function', default='gru', type=str, choices=['gru', 'lstm', 'layernorm_lstm', 'layernorm_gru'],
help='the update function to use.')
parser.add_argument('--wx_update', action='store_true', default=False,
help='The inputs for the update function considers the node feature of mlp.')
parser.add_argument('--no_keep_input', action='store_true', default=False,
help='no to use the input feature as the input to recurrent function.')
parser.add_argument('--aggr_state', action='store_true', default=False,
help='use the aggregated message as the previous state of recurrent function.')
parser.add_argument('--init_hidden', action='store_true',
default=False, help='whether to init the hidden state of node embeddings')
parser.add_argument('--num_rounds', type=int, default=1, metavar='N',
help='The number of rounds for grn propagation.'
'1 - the setting used in DAGNN/D-VAE')
parser.add_argument('--no_reverse', action='store_true', default=False,
help='Not to use the reverse layer to propagate the message.')
parser.add_argument('--seperate_hidden', action='store_true', default=False,
help='seperate node hidden states for forward layer and backward layer.')
parser.add_argument('--dim_hidden', type=int, default=64, metavar='N',
help='hidden size of recurrent unit.')
parser.add_argument('--dim_mlp', type=int, default=32, metavar='N',
help='hidden size of readout layers')
parser.add_argument('--dim_pred', type=int, default=1, metavar='N',
help='hidden size of readout layers')
parser.add_argument('--mul_mlp', action='store_true', default=False,
help='To use seperate MLP for different gate types.')
parser.add_argument('--wx_mlp', action='store_true', default=False,
help='The inputs for the mlp considers the node feature of mlp.')
# deepsat
parser.add_argument('--mask', action='store_true', default=False,
help='Use the mask for the node embedding or not')
# circuitsat/deepsat
parser.add_argument('--temperature', type=float, default=0.01,
help='initial value for temperature')
parser.add_argument('--eplison', type=float, default=0.4,
help='the anneling factore of temperature.')
parser.add_argument('--k_step', type=float, default=10.0,
help='the value for step funtion parameter k.')
parser.add_argument('--prob_loss', action='store_true', default=False,
help='To use the simulated probabilities as complementary supervision.')
parser.add_argument('--prob_weight', type=float, default=0.1,
help='the weight for simulated probability loss.')
# loss
parser.add_argument('--reg_loss', default='l1',
help='regression loss: sl1 | l1 | l2 | focalloss')
parser.add_argument('--cls_loss', default='bce_logit',
help='classification loss: bce - BCELoss | bce_logit - BCELossWithLogit | cross - CrossEntropyLoss')
parser.add_argument('--sat_loss', default='smoothstep', choices=['smoothstep'],
help='the loss for circuitsat: smoothstep')
# 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='30,45',
help='drop learning rate by 10.')
parser.add_argument('--grad_clip', type=float, default=0.,
help='gradiant clipping')
parser.add_argument('--num_epochs', type=int, default=60,
help='total training epochs.')
parser.add_argument('--batch_size', type=int, default=32,
help='batch size')
parser.add_argument('--master_batch_size', type=int, default=-1,
help='batch size on the master gpu.')
parser.add_argument('--num_iters', type=int, default=-1,
help='default: #samples / batch_size.')
parser.add_argument('--trainval_split', default=0.9, type=float,
help='the splitting setting for training dataset and validation dataset.')
parser.add_argument('--val_only', action='store_true',
help='Do the validation evaluation only.')
# test
parser.add_argument('--test_split', default='test', choices=['test', 'train', 'all'],
help='the split to use for testing.')
parser.add_argument('--cop_only', action='store_true',
help='only show the comparision between C1 and simluated probability.')
parser.add_argument('--test_num_rounds', default=10, type=int,
help='The number of rounds to be run during testing.')
#########################
# Dataset
#########################
parser.add_argument('--dataset_name', default='random_sr3_10_100', type=str)
args = parser.parse_args()
args.gpus_str = args.gpus
args.gpus = [int(gpu) for gpu in args.gpus.split(',')]
args.gpus = [i for i in range(len(args.gpus))] if args.gpus[0] >=0 else [-1]
args.lr_step = [int(i) for i in args.lr_step.split(',')]
# update data settings
DEFAULT_GATE_TO_INDEX = {"INPUT": 0, "AND": 1, "NOT": 2, "OR": 3, "NAND": 4, "NOR": 5, "XOR": 6}
args.gate_to_index = {}
if args.gate_types == '*':
args.gate_to_index = DEFAULT_GATE_TO_INDEX
else:
gate_types = args.gate_types.split(',')
for i in range(len(gate_types)):
args.gate_to_index[gate_types[i]] = i
args.num_gate_types = len(args.gate_to_index)
# check the relationship of `task`, `dataset` and `arch` comply with each other. TODO: optimize this part
if args.task == 'deepsat':
args.dataset = 'gate'
args.arch = 'deepsat'
elif args.task == 'neurosat':
args.dataset = 'cnf'
args.arch = 'neurosat'
else:
args.dataset = 'ckt'
args.arch = 'circuitsat'
args.reverse = not args.no_reverse
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.dataset_dir = os.path.join(args.root_dir, 'data', args.dataset_name)
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)
# TODO List
if args.mul_mlp: raise NotImplementedError
if args.debug == 2: raise NotImplementedError
if len(args.gpus) > 1: raise NotImplementedError('Only support single gpu right now.')
if args.seperate_hidden: raise NotImplementedError
if args.reg_loss == 'focalloss': raise NotImplementedError
return args