-
Notifications
You must be signed in to change notification settings - Fork 0
/
modelling.py
120 lines (92 loc) · 2.86 KB
/
modelling.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
"""
Script for constructing and training MDN model of transition density.
"""
import pickle
from pathlib import Path
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras import callbacks as cb
from preprocessing import Scaler # noqa: E402
tfkl = tf.keras.layers
tfpl = tfp.layers
tfd = tfp.distributions
kl = tfd.kullback_leibler
tf.keras.backend.set_floatx("float64")
# Model hyperparameters
N_C = 32
DT = 4
MODEL_DIR = (f"models/GDP_{DT:.0f}day_NC{N_C}/")
if not Path(MODEL_DIR).exists():
Path(MODEL_DIR).mkdir(parents=True)
# --- PREPARE DATA ---
DATA_DIR = f"data/GDP/{DT:.0f}day/"
X = np.load(DATA_DIR + "X0_train.npy")
Y = np.load(DATA_DIR + "DX_train.npy")
Xws = X.copy()
Xws[:, 0] -= 360.0
Xes = X.copy()
Xes[:, 0] += 360.0
# Periodicising X0.
X = np.concatenate((X, Xes, Xws), axis=0)
Y = np.concatenate((Y, Y, Y), axis=0)
Xscaler = Scaler(X)
Yscaler = Scaler(Y)
X_ = Xscaler.standardise(X)
Y_ = Yscaler.standardise(Y)
del X, Y
with open(MODEL_DIR + r"Xscaler.pkl", "wb") as file:
pickle.dump(Xscaler, file)
with open(MODEL_DIR + r"Yscaler.pkl", "wb") as file:
pickle.dump(Yscaler, file)
# --- BUILD MODEL ---
mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
model = tf.keras.Sequential(
[tfkl.Dense(256, activation='tanh'),
tfkl.Dense(256, activation='tanh'),
tfkl.Dense(256, activation='tanh'),
tfkl.Dense(256, activation='tanh'),
tfkl.Dense(512, activation='tanh'),
tfkl.Dense(512, activation='tanh'),
tfkl.Dense(N_C * 6, activation=None),
tfpl.MixtureSameFamily(N_C, tfpl.MultivariateNormalTriL(2))])
# --- TRAIN MODEL ---
LOG_FILE = "log.csv"
CHECKPOINT_FILE = "checkpoint_epoch_{epoch:02d}/weights"
TRAINED_FILE = "trained/weights"
# Training configuration
def nll(data_point, tf_distribution):
"""Negative log likelihood."""
return -tf_distribution.log_prob(data_point)
LOSS = nll
BATCH_SIZE = 8192
LEARNING_RATE = 5e-5
EPOCHS = 100000
OPTIMISER = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE)
VALIDATION_SPLIT = 0.2
# Callbacks
CSV_LOGGER = cb.CSVLogger(MODEL_DIR + LOG_FILE)
BATCHES_PER_EPOCH = int(
np.ceil(X_.shape[0] / BATCH_SIZE * (1 - VALIDATION_SPLIT)))
CHECKPOINTING = cb.ModelCheckpoint(
MODEL_DIR + CHECKPOINT_FILE,
save_freq=10 * BATCHES_PER_EPOCH,
verbose=1,
save_weights_only=True)
EARLY_STOPPING = cb.EarlyStopping(monitor="val_loss",
patience=50, min_delta=0.0)
CALLBACKS = [CHECKPOINTING, CSV_LOGGER, EARLY_STOPPING]
# Model compilation and training
model.compile(loss=LOSS, optimizer=OPTIMISER)
History = model.fit(
X_,
Y_,
epochs=EPOCHS,
callbacks=CALLBACKS,
batch_size=BATCH_SIZE,
shuffle=True,
validation_split=VALIDATION_SPLIT,
verbose=2,
)
model.save_weights(MODEL_DIR + TRAINED_FILE)