def train_add(func=lambda a, b: a + b, results_dir=None, reg_weight=5e-2, learning_rate=1e-2, n_epochs=10001):
"""Addition of two MNIST digits with a symbolic regression network."""
tf.reset_default_graph()
# Symbolic regression network to combine the conv net outputs
PRIMITIVE_FUNCS = [
*[functions.Constant()] * 2,
*[functions.Identity()] * 4,
*[functions.Square()] * 4,
*[functions.Sin()] * 2,
*[functions.Exp()] * 2,
*[functions.Sigmoid()] * 2,
*[functions.Product()] * 2,
]
sr_net = symbolic_network.SymbolicNet(2, funcs=PRIMITIVE_FUNCS, init_stddev=0.1) # Symbolic regression network
# Overall architecture
sym_digit_network = SymbolicDigit(sr_net=sr_net, normalize=normalize)
# Set up regularization term and training
penalty = regularization.l12_smooth(sr_net.get_weights())
penalty = reg_weight * penalty
sym_digit_network.set_training(reg=penalty)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True # Take up variable amount of memory on GPU
sess = tf.Session(config=config)
batch = batch_generator(batch_size=100)
# Train, and restart training if loss goes to NaN
loss_i = np.nan
while np.isnan(loss_i):
sess.run(tf.global_variables_initializer())
loss_i = sym_digit_network.train(sess, n_epochs, batch, func, lr_val=learning_rate)
if np.isnan(loss_i):
continue
# Freezing weights
sr_net = symbolic_network.MaskedSymbolicNet(sess, sr_net, threshold=0.01)
sym_digit_network = SymbolicDigitMasked(sym_digit_network, sr_net, normalize=normalize)
sym_digit_network.set_training()
# Training with frozen weights. Regularization is 0
loss_i = sym_digit_network.train(sess, n_epochs, batch, func, lr_val=learning_rate/10)
# Print out human-readable equation (with regularization)
weights = sess.run(sr_net.get_weights())
expr = pretty_print.network(weights, PRIMITIVE_FUNCS, ["z1", "z2"])
expr = normalize(expr)
print(expr)
# Calculate accuracy on test dataset
acc_test, error_test = sym_digit_network.calc_accuracy(X_test, y_test, func, sess)
result_str = 'Test accuracy: %g\n' % acc_test
print(result_str)
sym_digit_network.save_result(sess, results_dir, expr, result_str)
def __init__(self,
results_dir,
n_layers=2,
reg_weight=1e-2,
learning_rate=1e-2,
n_epochs1=10001):
super().__init__(
results_dir,
n_layers,
reg_weight,
learning_rate,
n_epochs1,
)
self.activation_funcs = [
[
*[functions.Constant()] * 2,
*[functions.Identity()] * 10,
# *[functions.Square()] * 4,
*[functions.Sin()] * 10,
# *[functions.Exp()] * 2,
# *[functions.Sigmoid()] * 2,
# *[functions.Reciprocal(1.0)] * 2,
*[functions.Product(1.0)] * 10,
],
[
*[functions.Constant()] * 2,
*[functions.Identity()] * 10,
# *[functions.Square()] * 4,
*[functions.Sin()] * 10,
# *[functions.Exp()] * 2,
# *[functions.Sigmoid()] * 2,
# *[functions.Division(1.0)] * 2,
*[functions.Product(1.0)] * 10,
],
[
# *[functions.Constant()] * 2,
# *[functions.Identity()] * 4,
# *[functions.Reciprocal(1.0)] * 2,
*[functions.Division(1.0)] * 2,
]
]
def __init__(self,
results_dir,
n_layers=2,
reg_weight=5e-3,
learning_rate=1e-2,
n_epochs1=10001,
n_epochs2=10001,
m=1):
"""Set hyper-parameters"""
self.activation_funcs = [
*[functions.Constant()] * 2 * m, *[functions.Identity()] * 4 * m,
*[functions.Square()] * 4 * m, *[functions.Sin()] * 2 * m,
*[functions.Exp()] * 2 * m, *[functions.Sigmoid()] * 2 * m,
*[functions.Product(1.0)] * 2 * m
]
self.n_layers = n_layers # Number of hidden layers
self.reg_weight = reg_weight # Regularization weight
self.learning_rate = learning_rate
self.summary_step = 1000 # Number of iterations at which to print to screen
self.n_epochs1 = n_epochs1
self.n_epochs2 = n_epochs2
self.m = m
if not os.path.exists(results_dir):
os.makedirs(results_dir)
self.results_dir = results_dir
# Save hyperparameters to file
result = {
"learning_rate": self.learning_rate,
"summary_step": self.summary_step,
"n_epochs1": self.n_epochs1,
"n_epochs2": self.n_epochs2,
"activation_funcs_name":
[func.name for func in self.activation_funcs],
"n_layers": self.n_layers,
"reg_weight": self.reg_weight,
'n_func_multiplier': self.m
}
with open(os.path.join(self.results_dir, 'params.pickle'), "wb+") as f:
pickle.dump(result, f)
def main(results_dir='results/sho/test', trials=1, learning_rate=1e-2, reg_weight=2e-4, timesteps=25, batch_size=129,
n_epochs1=2001, n_epochs2=5001, n_epochs3=5001):
# Hyperparameters
summary_step = 500
timesteps0 = 1
primitive_funcs = [
*[functions.Constant()] * 2,
*[functions.Identity()] * 4,
*[functions.Square()] * 4,
*[functions.Sin()] * 2,
*[functions.Exp()] * 2,
*[functions.Sigmoid()] * 2,
*[functions.Product(norm=0.1)] * 2,
]
# Import parabola data
data = np.load('dataset/sho.npz')
x_d = np.asarray(data["x_d"])
x_v = np.asarray(data["x_v"])
y_d = np.asarray(data["y_d"])
y_v = np.asarray(data["y_v"])
omega2_data = data["omega2"]
N = data["N"]
# Prepare data
x = np.stack((x_d, x_v), axis=2) # Shape (N, NT, 2)
y0 = np.stack((y_d[:, 0], y_v[:, 0]), axis=1) # Initial conditions for prediction y, fed into propagator
y_data = np.stack((y_d[:, 1:timesteps + 1], y_v[:, 1:timesteps + 1]), axis=2) # shape(NG, LENGTH, 2)
# Tensorflow placeholders for x, y0, y
x_input = tf.placeholder(shape=(None, x.shape[1], x.shape[2]), dtype=tf.float32, name="enc_input")
y0_input = tf.placeholder(shape=(None, 2), dtype=tf.float32, name="prop_input") # input is d, v
y_input = tf.placeholder(shape=(None, timesteps, 2), dtype=tf.float32, name="label_input")
length_input = tf.placeholder(dtype=tf.int32, shape=())
# Dynamics encoder
encoder = helpers.Encoder()
training = tf.placeholder_with_default(False, [])
z = encoder(x_input, training=training)
z_data = omega2_data[:, np.newaxis]
# Propagating decoders
prop_d = SymbolicNet(2, funcs=primitive_funcs)
prop_v = SymbolicNet(2, funcs=primitive_funcs)
prop_d.build(4)
prop_v.build(4)
# Building recurrent structure
rnn = tf.keras.layers.RNN(SymbolicCell(prop_d, prop_v), return_sequences=True)
y0_rnn = tf.concat([tf.expand_dims(y0_input, axis=1), tf.zeros((tf.shape(y0_input)[0], length_input - 1, 2))],
axis=1)
prop_input = tf.concat([y0_rnn, tf.keras.backend.repeat(z, length_input),
tf.ones((tf.shape(y0_input)[0], length_input, 1))], axis=2)
prop_output = rnn(prop_input)
epoch = tf.placeholder(tf.float32)
reg_freq = np.pi / (n_epochs1 + n_epochs2) / 1.1
reg_loss = tf.sin(reg_freq * epoch) ** 2 * regularization.l12_smooth(prop_d.get_weights()) + \
tf.sin(reg_freq * epoch) ** 2 * regularization.l12_smooth(prop_v.get_weights())
# reg_loss = regularization.l12_smooth(prop_d.get_weights()) + regularization.l12_smooth(prop_v.get_weights())
# Training
learning_rate_ph = tf.placeholder(tf.float32)
opt = tf.train.RMSPropOptimizer(learning_rate=learning_rate_ph)
reg_weight_ph = tf.placeholder(tf.float32)
error = tf.losses.mean_squared_error(labels=y_input[:, :length_input, :], predictions=prop_output)
loss = error + reg_weight_ph * reg_loss
train = tf.group([opt.minimize(loss), encoder.bn.updates])
batch = helpers.batch_generator([x, y_data, y0, z_data], N=N, batch_size=batch_size)
# Training session
with tf.Session() as sess:
for _ in range(trials):
loss_i = np.nan
while np.isnan(loss_i):
loss_list = []
error_list = []
reg_list = []
sess.run(tf.global_variables_initializer())
for i in range(n_epochs1 + n_epochs2):
if i < n_epochs1:
reg_weight_i = reg_weight / 5
learning_rate_i = learning_rate
length_i = min(i // 500 * 2 + timesteps0, timesteps)
else:
reg_weight_i = reg_weight
learning_rate_i = learning_rate / 5
length_i = timesteps
x_batch, y_batch, y0_batch, z_batch = next(batch)
feed_dict = {x_input: x_batch, y0_input: y0_batch, y_input: y_batch,
epoch: i, learning_rate_ph: learning_rate_i, training: True,
reg_weight_ph: reg_weight_i, length_input: length_i}
_ = sess.run(train, feed_dict=feed_dict)
if i % summary_step == 0 or i == n_epochs1 - 1:
feed_dict[training] = False
loss_i, error_i, reg_i = sess.run((loss, error, reg_loss), feed_dict=feed_dict)
z_arr = sess.run(z, feed_dict=feed_dict)
r = np.corrcoef(z_batch[:, 0], z_arr[:, 0])[1, 0]
loss_list.append(loss_i)
error_list.append(error_i)
reg_list.append(reg_i)
print("Epoch %d\tTotal loss: %f\tError: %f\tReg loss: %f\tCorrelation: %f"
% (i, loss_i, error_i, reg_i, r))
if np.isnan(loss_i):
break
# Setting small weights to 0 and freezing them
prop_d_masked = MaskedSymbolicNet(sess, prop_d, threshold=0.01)
prop_v_masked = MaskedSymbolicNet(sess, prop_v, threshold=0.01)
# Keep track of currently existing variables. When we rebuild the rnn, it makes new variables that we need
# to initialize. Later, we will use this to figure out what the uninitialized variables are.
temp = set(tf.global_variables())
# Rebuilding the decoding propagator. Remove regularization
rnn = tf.keras.layers.RNN(SymbolicCell(prop_d_masked, prop_v_masked), return_sequences=True)
prop_output = rnn(prop_input)
loss = tf.losses.mean_squared_error(labels=y_input[:, :length_input, :], predictions=prop_output)
train = tf.group([opt.minimize(loss), encoder.bn.updates])
weights_d = sess.run(prop_d_masked.get_weights())
expr_d = pretty_print.network(weights_d, primitive_funcs, ["d", "v", "z", 1])
print(expr_d)
weights_v = sess.run(prop_v_masked.get_weights())
expr_v = pretty_print.network(weights_v, primitive_funcs, ["d", "v", "z", 1])
print(expr_v)
print("Frozen weights. Next stage of training.")
# Initialize only the uninitialized variables.
sess.run(tf.variables_initializer(set(tf.global_variables()) - temp))
for i in range(n_epochs3):
x_batch, y_batch, y0_batch, z_batch = next(batch)
feed_dict = {x_input: x_batch, y0_input: y0_batch, y_input: y_batch,
epoch: 0, learning_rate_ph: learning_rate / 10, training: True, reg_weight_ph: 0,
length_input: length_i}
_ = sess.run(train, feed_dict=feed_dict)
if i % summary_step == 0:
feed_dict[training] = False
loss_i, error_i, reg_i = sess.run((loss, error, reg_loss), feed_dict=feed_dict)
z_arr = sess.run(z, feed_dict=feed_dict)
r = np.corrcoef(z_batch[:, 0], z_arr[:, 0])[1, 0]
loss_list.append(loss_i)
error_list.append(error_i)
reg_list.append(reg_i)
print("Epoch %d\tError: %g\tCorrelation: %f" % (i, error_i, r))
weights_d = sess.run(prop_d_masked.get_weights())
expr_d = pretty_print.network(weights_d, primitive_funcs, ["d", "v", "z", 1])
print(expr_d)
weights_v = sess.run(prop_v_masked.get_weights())
expr_v = pretty_print.network(weights_v, primitive_funcs, ["d", "v", "z", 1])
print(expr_v)
# Save results
results = {
"summary_step": summary_step,
"learning_rate": learning_rate,
"n_epochs1": n_epochs1,
"n_epochs2": n_epochs2,
"reg_weight": reg_weight,
"timesteps": timesteps,
"timesteps0": timesteps0,
"weights_d": weights_d,
"weights_v": weights_v,
"loss_plot": loss_list,
"error_plot": error_list,
"reg_plot": reg_list,
"expr_d": expr_d,
"expr_v": expr_v
}
trial_dir = helpers.get_trial_path(results_dir) # Get directory in which to save trial results
tf.saved_model.simple_save(sess, trial_dir,
inputs={"x": x_input, "y0": y0_input, "training": training},
outputs={"z": z, "y": prop_output})
# Save a summary of the parameters and results
with open(os.path.join(trial_dir, 'summary.pickle'), "wb+") as f:
pickle.dump(results, f)
with open(os.path.join(results_dir, 'eq_summary.txt'), 'a') as f:
f.write(str(expr_d) + "\n")
f.write(str(expr_v) + "\n")
f.write("Error: %f\n\n" % error_list[-1])
def main(results_dir='results/kinematics/test',
learning_rate=1e-2,
reg_weight=1e-3,
n_epochs1=5001,
n_epochs2=5001,
timesteps=5):
# Hyperparameters
summary_step = 500
timesteps0 = 1
# Import kinematics data
data = np.load('dataset/kinematic.npz')
x_d = np.asarray(data["x_d"])
x_v = np.asarray(data["x_v"])
y_d = np.asarray(data["y_d"])
y_v = np.asarray(data["y_v"])
a_data = np.asarray(data["g"])
# Prepare data
# The first few time steps are reserved for the symbolic regression propagator
x = np.stack((x_d, x_v), axis=2) # Shape (N, NT, 2)
y0 = np.stack((y_d[:, 0], y_v[:, 0]),
axis=1) # Input into the symbolic propagator
label_data = np.stack((y_d[:, 1:timesteps + 1], y_v[:, 1:timesteps + 1]),
axis=2) # shape(NG, timesteps, 2)
# Encoder
encoder = helpers.Encoder() # layer should end with 1, which is the output
x_input = tf.placeholder(shape=(None, x.shape[1], x.shape[2]),
dtype=tf.float32,
name="enc_input")
y_input = tf.placeholder(shape=(None, timesteps, 2),
dtype=tf.float32,
name="label_input")
training = tf.placeholder_with_default(False, [])
z = encoder(x_input, training=training)
# z = np.array(a_data)[:, np.newaxis] # uncomment to ignore the autoencoder
# Propagating decoder
primitive_funcs = [
*[functions.Constant()] * 2,
*[functions.Identity()] * 4,
*[functions.Square()] * 4,
*[functions.Sin()] * 2,
*[functions.Exp()] * 2,
*[functions.Sigmoid()] * 2,
*[functions.Product(norm=0.1)] * 2,
]
prop_d = SymbolicNet(2, funcs=primitive_funcs)
prop_v = SymbolicNet(2, funcs=primitive_funcs)
prop_input = tf.placeholder(shape=(None, 2),
dtype=tf.float32,
name="prop_input") # input is d, v
def rec_sr(y0_input, enc_output, length, prop1=prop_d, prop2=prop_v):
rec_input = [y0_input]
for i in range(length):
full_input = tf.concat(
[rec_input[i], enc_output,
tf.ones_like(enc_output)],
axis=1,
name="full_input") # d, v, z
rec_input.append(
tf.concat(
[prop1(full_input), prop2(full_input)],
axis=1,
name="c_prop_input"))
output = tf.stack(rec_input[1:], axis=1) # Ignore initial conditions
return output
y_hat_start = rec_sr(prop_input, z, timesteps0, prop_d, prop_v)
y_hat_full = rec_sr(prop_input, z, timesteps, prop_d, prop_v)
# Label and errors
epoch = tf.placeholder(tf.float32)
reg_weight_ph = tf.placeholder(tf.float32)
reg_loss = regularization.l12_smooth(
prop_d.get_weights()) + regularization.l12_smooth(prop_v.get_weights())
# Training
learning_rate_ph = tf.placeholder(tf.float32)
opt = tf.train.RMSPropOptimizer(learning_rate=learning_rate_ph)
def define_loss(prop_output, length):
error = tf.losses.mean_squared_error(
labels=y_input[:, :length, :],
predictions=prop_output[:, :length, :])
loss = error + reg_weight_ph * reg_loss
train = opt.minimize(loss)
train = tf.group([train, encoder.bn.updates])
return error, loss, train
error_start, loss_start, train_start = define_loss(y_hat_start, timesteps0)
error_full, loss_full, train_full = define_loss(y_hat_full, timesteps)
# Training session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True # Take up variable amount of memory on GPU
with tf.Session(config=config) as sess:
loss_i = np.nan
while np.isnan(loss_i):
loss_list = []
error_list = []
reg_list = []
error, loss, train = error_start, loss_start, train_start
sess.run(tf.global_variables_initializer())
for i in range(n_epochs1):
feed_dict = {
x_input: x,
prop_input: y0,
y_input: label_data,
epoch: 0,
learning_rate_ph: learning_rate,
training: True,
reg_weight_ph: reg_weight
}
_ = sess.run(train, feed_dict=feed_dict)
if i % summary_step == 0:
feed_dict[training] = False
print_loss, print_error, print_l12 = sess.run(
(loss, error, reg_loss), feed_dict=feed_dict)
loss_list.append(print_loss)
error_list.append(print_error)
reg_list.append(print_l12)
print("Epoch %d\tTotal loss: %f\tError: %f\tReg loss: %f" %
(i, print_loss, print_error, print_l12))
loss_i = print_loss
if i > 2000:
error, loss, train = error_full, loss_full, train_full
if np.isnan(loss_i):
break
# Setting small weights to 0 and freezing them
prop_d_masked = MaskedSymbolicNet(sess, prop_d, threshold=0.1)
prop_v_masked = MaskedSymbolicNet(sess, prop_v, threshold=0.1)
# Rebuilding the decoding propagator
prop_output_masked = rec_sr(prop_input, z, timesteps, prop_d_masked,
prop_v_masked)
error, loss, train = define_loss(prop_output_masked, timesteps)
weights_d = sess.run(prop_d_masked.get_weights())
expr_d = pretty_print.network(weights_d, primitive_funcs,
["d", "v", "z", 1])
print(expr_d)
weights_v = sess.run(prop_v_masked.get_weights())
expr_v = pretty_print.network(weights_v, primitive_funcs,
["d", "v", "z", 1])
print(expr_v)
print("Frozen weights. Next stage of training.")
for i in range(n_epochs2):
feed_dict = {
x_input: x,
prop_input: y0,
y_input: label_data,
epoch: 0,
learning_rate_ph: learning_rate / 10,
training: True,
reg_weight_ph: 0
}
_ = sess.run(train, feed_dict=feed_dict)
if i % summary_step == 0:
feed_dict[training] = False
print_loss, print_error, print_l12 = sess.run(
(loss, error, reg_loss), feed_dict=feed_dict)
loss_list.append(print_loss)
error_list.append(print_error)
reg_list.append(print_l12)
print("Epoch %d\tError: %g" % (i, print_error))
weights_d = sess.run(prop_d_masked.get_weights())
expr_d = pretty_print.network(weights_d, primitive_funcs,
["d", "v", "z", 1])
print(expr_d)
weights_v = sess.run(prop_v_masked.get_weights())
expr_v = pretty_print.network(weights_v, primitive_funcs,
["d", "v", "z", 1])
print(expr_v)
# Save results
results = {
"timesteps": timesteps,
"summary_step": summary_step,
"learning_rate": learning_rate,
"n_epochs1": n_epochs1,
"n_epochs2": n_epochs2,
"reg_weight_ph": reg_weight,
"weights_d": weights_d,
"weights_v": weights_v,
"loss_plot": loss_list,
"error_plot": error_list,
"l12_plot": reg_list,
"expr_d": expr_d,
"expr_v": expr_v
}
trial_dir = helpers.get_trial_path(
results_dir) # Get directory in which to save trial results
tf.saved_model.simple_save(sess,
trial_dir,
inputs={
"x": x_input,
"y0": prop_input,
"training": training
},
outputs={
"z": z,
"y": y_hat_full
})
# Save a summary of the parameters and results
with open(os.path.join(trial_dir, 'summary.pickle'), "wb+") as f:
pickle.dump(results, f)
def main(results_dir='results/kinematics/test', learning_rate=1e-2, reg_weight=1e-3, n_epochs=10001,
timesteps=5):
tf.reset_default_graph()
# Hyperparameters
summary_step = 1000
# tf.set_random_seed(0)
# Import parabola data
data = np.load('dataset/kinematic.npz')
x_d = np.asarray(data["x_d"])
x_v = np.asarray(data["x_v"])
y_d = np.asarray(data["y_d"])
y_v = np.asarray(data["y_v"])
a_data = np.asarray(data["g"])
# Prepare data
# The first few time steps are reserved for the symbolic regression propagator
x = np.stack((x_d, x_v), axis=2) # Shape (N, NT, 2)
y0 = np.stack((y_d[:, 0], y_v[:, 0]), axis=1) # Input into the symbolic propagator
y_data = np.stack((y_d[:, 1:timesteps + 1], y_v[:, 1:timesteps + 1]), axis=2) # shape(NG, LENGTH, 2)
# Encoder
encoder = helpers.Encoder() # layer should end with 1, which is the output
x_input = tf.placeholder(shape=(None, x.shape[1], x.shape[2]), dtype=tf.float32, name="enc_input")
y_input = tf.placeholder(shape=(None, timesteps, 2), dtype=tf.float32, name="label_input")
y0_input = tf.placeholder(shape=(None, 2), dtype=tf.float32, name="y_input") # input is d, v
length_input = tf.placeholder(dtype=tf.int32, shape=())
training = tf.placeholder_with_default(False, [])
z = encoder(x_input, training=training)
# enc_output = np.array(g_data)[:, np.newaxis] # uncomment to ignore the autoencoder
# Build EQL network for the propagating decoder
primitive_funcs = [
*[functions.Constant()] * 2,
*[functions.Identity()] * 4,
*[functions.Square()] * 4,
*[functions.Sin()] * 2,
*[functions.Exp()] * 2,
*[functions.Sigmoid()] * 2,
*[functions.Product(norm=0.1)] * 2,
]
prop_d = SymbolicNetL0(2, funcs=primitive_funcs)
prop_v = SymbolicNetL0(2, funcs=primitive_funcs)
prop_d.build(4)
prop_v.build(4)
# Build recurrent structure
rnn = tf.keras.layers.RNN(SymbolicCell(prop_d, prop_v), return_sequences=True)
y0_rnn = tf.concat([tf.expand_dims(y0_input, axis=1), tf.zeros((tf.shape(y0_input)[0], length_input - 1, 2))], axis=1)
prop_input = tf.concat([y0_rnn, tf.keras.backend.repeat(z, length_input),
tf.ones((tf.shape(y0_input)[0], length_input, 1))], axis=2)
y_hat = rnn(prop_input)
# Label and errors
reg_loss = prop_d.get_loss() + prop_v.get_loss()
# Training
learning_rate_ph = tf.placeholder(tf.float32)
opt = tf.train.RMSPropOptimizer(learning_rate=learning_rate_ph)
error = tf.losses.mean_squared_error(labels=y_input[:, :length_input, :], predictions=y_hat)
loss = error + reg_weight * reg_loss
train = opt.minimize(loss)
train = tf.group([train, encoder.bn.updates])
# Training session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True # Take up variable amount of memory on GPU
with tf.Session(config=config) as sess:
loss_i = np.nan
while np.isnan(loss_i):
loss_list = []
error_list = []
reg_list = []
sess.run(tf.global_variables_initializer())
length_i = 1
for i in range(n_epochs):
lr_i = learning_rate
feed_dict = {x_input: x, y0_input: y0, y_input: y_data,
learning_rate_ph: lr_i, training: True, length_input: length_i}
_ = sess.run(train, feed_dict=feed_dict)
if i % summary_step == 0:
feed_dict[training] = False
loss_val, error_val, reg_val = sess.run((loss, error, reg_loss), feed_dict=feed_dict)
loss_list.append(loss_val)
error_list.append(error_val)
reg_list.append(reg_val)
print("Epoch %d\tTotal loss: %f\tError: %f\tReg loss: %f" % (i, loss_val, error_val, reg_val))
loss_i = loss_val
if i > 3000:
length_i = timesteps
if np.isnan(loss_i):
break
weights_d = sess.run(prop_d.get_weights())
expr_d = pretty_print.network(weights_d, primitive_funcs, ["d", "v", "z", 1])
print(expr_d)
weights_v = sess.run(prop_v.get_weights())
expr_v = pretty_print.network(weights_v, primitive_funcs, ["d", "v", "z", 1])
print(expr_v)
# z_arr = sess.run(enc_output, feed_dict=feed_dict)
# Save results
results = {
"timesteps": timesteps,
"summary_step": summary_step,
"learning_rate": learning_rate,
"N_EPOCHS": n_epochs,
"reg_weight": reg_weight,
"weights_d": weights_d,
"weights_v": weights_v,
"loss_plot": loss_list,
"error_plot": error_list,
"l12_plot": reg_list,
"expr_d": expr_d,
"expr_v": expr_v
}
trial_dir = helpers.get_trial_path(results_dir) # Get directory in which to save trial results
tf.saved_model.simple_save(sess, trial_dir,
inputs={"x": x_input, "y0": y0_input, "training": training},
outputs={"z": z, "y": y_hat})
# Save a summary of the parameters and results
with open(os.path.join(trial_dir, 'summary.pickle'), "wb+") as f:
pickle.dump(results, f)
def main(results_dir='results/sho/test',
trials=20,
learning_rate=1e-3,
reg_weight=1e-3,
timesteps=25,
batch_size=128,
n_epochs1=10001,
n_epochs2=10001):
# Hyperparameters
summary_step = 1000
primitive_funcs = [
*[functions.Constant()] * 2,
*[functions.Identity()] * 4,
*[functions.Square()] * 4,
*[functions.Sin()] * 2,
*[functions.Exp()] * 2,
*[functions.Sigmoid()] * 2,
*[functions.Product(norm=0.1)] * 2,
]
# Import parabola data
data = np.load('dataset/sho.npz')
x_d = np.asarray(data["x_d"])
x_v = np.asarray(data["x_v"])
y_d = np.asarray(data["y_d"])
y_v = np.asarray(data["y_v"])
omega2_data = data["omega2"]
N = data["N"]
# Prepare data
x = np.stack((x_d, x_v), axis=2) # Shape (N, NT, 2)
y0 = np.stack(
(y_d[:, 0], y_v[:, 0]),
axis=1) # Initial conditions for prediction y, fed into propagator
y_data = np.stack((y_d[:, 1:timesteps + 1], y_v[:, 1:timesteps + 1]),
axis=2) # shape(NG, timesteps, 2)
z_data = omega2_data[:, np.newaxis]
# Tensorflow placeholders for x, y0, y
x_input = tf.placeholder(shape=(None, x.shape[1], x.shape[2]),
dtype=tf.float32,
name="enc_input")
y0_input = tf.placeholder(shape=(None, 2),
dtype=tf.float32,
name="prop_input") # input is d, v
y_input = tf.placeholder(shape=(None, timesteps, 2),
dtype=tf.float32,
name="label_input")
length_input = tf.placeholder(dtype=tf.int32, shape=())
# Dynamics encoder
encoder = helpers.Encoder(n_filters=[16, 16, 16, 16])
training = tf.placeholder_with_default(False, [])
z = encoder(x_input, training=training)
# Propagating decoders
prop_d = SymbolicNetL0(2, funcs=primitive_funcs)
prop_v = SymbolicNetL0(2, funcs=primitive_funcs)
prop_d.build(4)
prop_v.build(4)
# Building recurrent structure
rnn = tf.keras.layers.RNN(SymbolicCell(prop_d, prop_v),
return_sequences=True)
y0_rnn = tf.concat([
tf.expand_dims(y0_input, axis=1),
tf.zeros((tf.shape(y0_input)[0], length_input - 1, 2))
],
axis=1)
prop_input = tf.concat([
y0_rnn,
tf.keras.backend.repeat(z, length_input),
tf.ones((tf.shape(y0_input)[0], length_input, 1))
],
axis=2)
y_hat = rnn(prop_input)
length_list = [1, 2, 3, 4, 5, 7, 10, 15,
25] # Slowly increase the length of propagation
# Training
learning_rate_ph = tf.placeholder(tf.float32)
opt = tf.train.RMSPropOptimizer(learning_rate=learning_rate_ph)
reg_weight_ph = tf.placeholder(tf.float32)
reg_loss = prop_d.get_loss() + prop_v.get_loss()
error = tf.losses.mean_squared_error(labels=y_input[:, :length_input, :],
predictions=y_hat)
loss = error + reg_weight_ph * reg_loss
train = tf.group([opt.minimize(loss), encoder.bn.updates])
batch = helpers.batch_generator([x, y_data, y0, z_data],
N=N,
batch_size=batch_size)
# Training session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
for _ in range(trials):
loss_i = np.nan
while np.isnan(loss_i):
loss_list = []
error_list = []
reg_list = []
sess.run(tf.global_variables_initializer())
length_i = 1
for i in range(n_epochs1 + n_epochs2):
if i < n_epochs1:
lr_i = learning_rate
else:
lr_i = learning_rate / 10
x_batch, y_batch, y0_batch, z_batch = next(batch)
feed_dict = {
x_input: x_batch,
y0_input: y0_batch,
y_input: y_batch,
learning_rate_ph: lr_i,
training: True,
reg_weight_ph: reg_weight,
length_input: length_i
}
_ = sess.run(train, feed_dict=feed_dict)
if i % summary_step == 0:
feed_dict[training] = False
loss_i, error_i, reg_i, z_arr = sess.run(
(loss, error, reg_loss, z), feed_dict=feed_dict)
r = np.corrcoef(z_batch[:, 0], z_arr[:, 0])[1, 0]
loss_list.append(loss_i)
error_list.append(error_i)
reg_list.append(reg_i)
print(
"Epoch %d\tTotal loss: %f\tError: %f\tReg loss: %f\tCorrelation: %f"
% (i, loss_i, error_i, reg_i, r))
if np.isnan(loss_i):
break
i_length = min(i // 1000, len(length_list) - 1)
length_i = length_list[i_length]
weights_d = sess.run(prop_d.get_weights())
expr_d = pretty_print.network(weights_d, primitive_funcs,
["d", "v", "z", 1])
print(expr_d)
weights_v = sess.run(prop_v.get_weights())
expr_v = pretty_print.network(weights_v, primitive_funcs,
["d", "v", "z", 1])
print(expr_v)
print("Done. Saving results.")
# z_arr = sess.run(z, feed_dict=feed_dict)
# Save results
results = {
"summary_step": summary_step,
"learning_rate": learning_rate,
"n_epochs1": n_epochs1,
"reg_weight": reg_weight,
"timesteps": timesteps,
"weights_d": weights_d,
"weights_v": weights_v,
"loss_plot": loss_list,
"error_plot": error_list,
"reg_plot": reg_list,
"expr_d": expr_d,
"expr_v": expr_v
}
trial_dir = helpers.get_trial_path(
results_dir) # Get directory in which to save trial results
tf.saved_model.simple_save(sess,
trial_dir,
inputs={
"x": x_input,
"y0": y0_input,
"training": training
},
outputs={
"z": z,
"y": y_hat
})
# Save a summary of the parameters and results
with open(os.path.join(trial_dir, 'summary.pickle'), "wb+") as f:
pickle.dump(results, f)
with open(os.path.join(results_dir, 'eq_summary.txt'), 'a') as f:
f.write(str(expr_d) + "\n")
f.write(str(expr_v) + "\n")
f.write("Error: %f\n\n" % error_list[-1])