-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
103 lines (85 loc) · 4.09 KB
/
main.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
import os
import numpy as np
import tensorflow as tf
import random
from train import run
import warnings
warnings.filterwarnings('ignore')
FLAGS = tf.compat.v1.app.flags.FLAGS
"""Set up dataset."""
tf.compat.v1.app.flags.DEFINE_string(
'dataset', 'sin', 'dataset/scenario: sin, step, abs, linear, poly2d, poly3d, mnistz, mnistc, mnistzc')
tf.compat.v1.app.flags.DEFINE_integer('num_train', 500, 'training data')
tf.compat.v1.app.flags.DEFINE_integer('num_valid', 500, 'validation data')
tf.compat.v1.app.flags.DEFINE_integer('num_test', 500, 'test data')
"""Set up training."""
tf.compat.v1.app.flags.DEFINE_integer(
'rep_net_layer', 2, 'layers of representation networks')
tf.compat.v1.app.flags.DEFINE_integer(
'x_net_layer', 2, 'layers of treatment prediction network')
tf.compat.v1.app.flags.DEFINE_integer(
'emb_net_layer', 2, 'layers of embedding network')
tf.compat.v1.app.flags.DEFINE_integer(
'y_net_layer', 2, 'layers of outcome prediction network')
tf.compat.v1.app.flags.DEFINE_integer('emb_dim', 4, 'embedding dimension')
tf.compat.v1.app.flags.DEFINE_integer('rep_dim', 4, 'representation dimension')
tf.compat.v1.app.flags.DEFINE_float(
'lrate', 1e-3, 'learning rate of optimizer')
tf.compat.v1.app.flags.DEFINE_float('dropout', 0., 'drop rate of dropout')
tf.compat.v1.app.flags.DEFINE_integer(
'epochs', 1000, 'training epochs of AutoIV')
tf.compat.v1.app.flags.DEFINE_integer(
'opt_lld_step', 1, 'steps of likelihood optimizer')
tf.compat.v1.app.flags.DEFINE_integer(
'opt_bound_step', 1, 'steps of bound optimizer')
tf.compat.v1.app.flags.DEFINE_integer(
'opt_2stage_step', 1, 'steps of two-stage optimizer')
tf.compat.v1.app.flags.DEFINE_float(
'sigma', 0.1, 'hyper-parameter sigma of RBF kernel')
tf.compat.v1.app.flags.DEFINE_integer(
'interval', 2, 'print times during training')
tf.compat.v1.app.flags.DEFINE_integer('exp_num', 3, 'experiment runs')
""" Set up experiments. """
tf.compat.v1.app.flags.DEFINE_boolean(
'gen_new', False, 'whether generate new data')
tf.compat.v1.app.flags.DEFINE_boolean(
'del_res', False, 'whether delete all the previous results')
tf.compat.v1.app.flags.DEFINE_string(
'res_path', 'AutoIV-results/', 'result path')
tf.compat.v1.app.flags.DEFINE_string(
'res_file', 'summary.csv', 'result summary csv file')
tf.compat.v1.app.flags.DEFINE_string('gpu', '0', 'which GPU to use')
tf.compat.v1.app.flags.DEFINE_integer('seed', 0, 'seed')
os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.gpu
seed = FLAGS.seed
random.seed(seed)
np.random.seed(seed)
tf.set_random_seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
def main_run(FLAGS):
"""Set hyper-parameters."""
coefs = {'coef_cx2y': 1, 'coef_zc2x': 1, 'coef_lld_zx': 1,
'coef_lld_zy': 1, 'coef_lld_cx': 1, 'coef_lld_cy': 1,
'coef_lld_zc': 1, 'coef_bound_zx': 1, 'coef_bound_zy': 1,
'coef_bound_cx': 1, 'coef_bound_cy': 1, 'coef_bound_zc': 1, 'coef_reg': 0.001}
num = {'train': FLAGS.num_train,
'valid': FLAGS.num_valid, 'test': FLAGS.num_test}
layers = {'rep_net_layer': FLAGS.rep_net_layer, 'x_net_layer': FLAGS.x_net_layer,
'emb_net_layer': FLAGS.emb_net_layer, 'y_net_layer': FLAGS.y_net_layer}
opt_steps = {'opt_lld_step': FLAGS.opt_lld_step, 'opt_bound_step': FLAGS.opt_bound_step,
'opt_2stage_step': FLAGS.opt_2stage_step}
data_path = 'data/{}/{}-train_{}/'.format(
FLAGS.dataset, FLAGS.dataset, FLAGS.num_train)
all_paras = {'dataset': FLAGS.dataset, 'data_path': data_path, 'coefs': coefs,
'rep_dim': FLAGS.rep_dim, 'lrate': FLAGS.lrate, 'dropout': FLAGS.dropout,
'emb_dim': FLAGS.emb_dim, 'epochs': FLAGS.epochs, 'interval': FLAGS.interval,
'exp_num': FLAGS.exp_num, 'gen_new': FLAGS.gen_new, 'del_res': FLAGS.del_res,
'sigma': FLAGS.sigma, 'res_path': FLAGS.res_path, 'visible_gpu': FLAGS.gpu}
all_paras.update(num)
all_paras.update(layers)
all_paras.update(opt_steps)
print('\n\n' + '=' * 50)
print('Run experiment.\n')
_, _ = run(all_paras)
if __name__ == '__main__':
main_run(FLAGS)