-
Notifications
You must be signed in to change notification settings - Fork 6
/
train_network.py
112 lines (92 loc) · 5.08 KB
/
train_network.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
from data_gen import pendulum
from sklearn.model_selection import train_test_split
from preprocessing import scale
import pickle
from models.mlp_tf import mlp, mlp_flipout
from models.cd import make_model
import argparse
import os
import numpy as np
parser = argparse.ArgumentParser()
parser.add_argument('--model-type', default='de', type=str, help='model type')
parser.add_argument('--t-spread-min', default=0.01, type=float, help='minimum T spread (rel noise on measurements)')
parser.add_argument('--t-spread-max', default=0.1, type=float, help='maximum T spread (rel noise on measurements)')
parser.add_argument('--ell-spread-min', default=0.02, type=float, help='minimum L spread (rel noise on measurements)')
parser.add_argument('--ell-spread-max', default=0.02, type=float, help='maximum L spread (rel noise on measurements)')
parser.add_argument('--n', default=100000, type=int, help='number of training and validation examples')
parser.add_argument('--n-test', default=10000, type=int, help='number of test examples')
parser.add_argument('--n-epochs', default=100, type=int, help='number of training epochs')
parser.add_argument('--data-dir', default='data/', type=str, help='data directory')
val_proportion = 0.1
n_models = 10
n_neurons = 100
def main(model_type, t_spread_min, t_spread_max, ell_spread_min, ell_spread_max, n, n_test, n_epochs, data_dir):
# Generate data
feat, y, _, _ = pendulum(n=n, t_spread=[t_spread_min, t_spread_max], ell_spread=[ell_spread_min, ell_spread_max])
# Set up data
x_train, x_val, y_train, y_val = train_test_split(feat, y, test_size=val_proportion, random_state=42)
x_scaler, x_train, x_val = scale(x_train, x_val)
y_scaler, y_train, y_val = scale(y_train, y_val)
t_range_str = f'trange{int(100*t_spread_min)}to{int(100*t_spread_max)}'
model_name = f'{model_type}_{t_range_str}_{n_epochs}ep'
os.makedirs(data_dir, exist_ok=True)
if not os.path.isfile(f'{data_dir}x_scaler_{t_range_str}.pkl'):
with open(f'{data_dir}x_scaler_{t_range_str}.pkl', 'wb') as file_pi:
pickle.dump(x_scaler, file_pi)
with open(f'{data_dir}y_scaler_{t_range_str}.pkl', 'wb') as file_pi:
pickle.dump(y_scaler, file_pi)
# train and save models
model_number = 1
while os.path.isfile(f'{data_dir}model_{model_name}_{str(model_number).zfill(3)}.h5'):
model_number += 1
if model_type == 'de':
models = [mlp(loss='nll') for _ in range(n_models)]
elif model_type == 'cd':
n_features = x_train.shape[1]
n_outputs = y_train.shape[1]
dropout_reg = 2. / n
models = [make_model(n_features, n_outputs, n_neurons, dropout_reg)]
elif model_type == 'bnn':
models = [mlp_flipout()]
else:
raise ValueError(f'Model type {model_type} not recognized!')
for j, mod in enumerate(models):
print(f'Model {j+1}')
history = mod.fit(x_train, y_train, epochs=n_epochs, validation_data=(x_val, y_val))
mod.save_weights(f'{data_dir}model_{model_name}_{str(model_number+j).zfill(3)}.h5')
with open(f'{data_dir}history_{model_name}_{str(model_number+j).zfill(3)}.pkl', 'wb') as file_pi:
pickle.dump(history.history, file_pi)
# Generate test set
feat_test, _, _, _ = pendulum(n=n_test, t_spread=[t_spread_min, t_spread_max],
ell_spread=[ell_spread_min, ell_spread_max], seed=666)
feat_test = x_scaler.transform(feat_test)
# make predictions
if model_type == 'de':
y_pred = []
for model in models:
y_pred.append(model(feat_test.astype('float32')))
elif model_type == 'cd':
y_pred = np.array([models[0].predict(feat_test) for _ in range(n_models)])
elif model_type == 'bnn':
y_pred = [models[0](feat_test.astype('float32')) for _ in range(n_models)]
if model_type == 'de' or model_type == 'bnn':
y_pred_val = [pred.loc.numpy() for pred in y_pred]
y_pred_unc = [pred.scale.numpy() for pred in y_pred]
elif model_type == 'cd':
y_pred_val = y_pred[:, :, :1]
y_pred_unc = np.sqrt(np.exp(y_pred[:, :, 1:]))
y_pred_val_resc = [y_scaler.inverse_transform(y) for y in y_pred_val]
y_pred_unc_resc = [y / y_scaler.scale_[0] for y in y_pred_unc]
y_pred_val_resc = np.array(y_pred_val_resc).reshape((n_models, n_test))
y_pred_unc_resc = np.array(y_pred_unc_resc).reshape((n_models, n_test))
y_pred_mean = np.mean(y_pred_val_resc, axis=0)
y_pred_ep_unc = np.std(y_pred_val_resc, axis=0)
y_pred_al_unc = np.sqrt(np.mean(y_pred_unc_resc * y_pred_unc_resc, axis=0))
y_pred_unc = np.sqrt(y_pred_al_unc ** 2 + y_pred_ep_unc ** 2)
np.save(f'{data_dir}y_pred_test_{model_name}_{str(model_number).zfill(3)}.npy', y_pred_mean)
np.save(f'{data_dir}y_pred_test_alunc_{model_name}_{str(model_number).zfill(3)}.npy', y_pred_al_unc)
np.save(f'{data_dir}y_pred_test_epunc_{model_name}_{str(model_number).zfill(3)}.npy', y_pred_ep_unc)
np.save(f'{data_dir}y_pred_test_prunc_{model_name}_{str(model_number).zfill(3)}.npy', y_pred_unc)
if __name__ == '__main__':
args = parser.parse_args()
main(**vars(args))