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 forward(self, input, target, a=0.05):
"""input: predictions"""
error = F.mse_loss(input, target)
# error_test = tf.losses.mean_squared_error(labels=y_test, predictions=y_hat)
# TODO: make l12_smooth a Module
# TODO: regularization
weights = sym.get_weights()
# reg_loss = l12_smooth()
if type(weights) == list:
reg_loss = sum([l12_smooth(tensor) for tensor in weights])
else:
smooth_abs = torch.where(
torch.abs(weights) < a,
torch.pow(weights, 4) / (-8 * a**3) +
torch.square(weights) * 3 / 4 / a + 3 * a / 8,
torch.abs(weights))
reg_loss = torch.sum(torch.sqrt(smooth_abs))
loss = error + self.reg_weight * reg_loss
return loss
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 train_add_test(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.
Withold sums > 15 for test data"""
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) # 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())
epoch = tf.placeholder_with_default(0.0, [])
penalty = tf.sin(np.pi / n_epochs / 1.1 * epoch) ** 2 * 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)
def train_fun(y):
return y < 15
def test_fun(y):
return np.logical_not(train_fun(y))
# 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, epoch, lr_val=learning_rate, train_fun=train_fun)
if np.isnan(loss_i):
continue
# Freezing weights
sr_net_masked = symbolic_network.MaskedSymbolicNet(sess, sr_net, threshold=0.01)
sym_digit_network = SymbolicDigitMasked(sym_digit_network, sr_net_masked, normalize=normalize)
sym_digit_network.set_training()
loss_i = sym_digit_network.train(sess, n_epochs, batch, func, lr_val=learning_rate/10, train_fun=train_fun)
# 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_train, error_train = sym_digit_network.calc_accuracy(X_train, y_train, func, sess)
acc_train1, error_train1 = sym_digit_network.calc_accuracy(X_train, y_train, func, sess, filter_fun=train_fun)
acc_train2, error_train2 = sym_digit_network.calc_accuracy(X_train, y_train, func, sess, filter_fun=test_fun)
acc_test, error_test = sym_digit_network.calc_accuracy(X_test, y_test, func, sess)
acc_test1, error_test1 = sym_digit_network.calc_accuracy(X_test, y_test, func, sess, filter_fun=train_fun)
acc_test2, error_test2 = sym_digit_network.calc_accuracy(X_test, y_test, func, sess, filter_fun=test_fun)
result_str = "Train digits overall accuracy: %.3f\ttrain sum accuracy: %.3f\t test sum accuracy: %.3f\n" \
"Train digits overall error: %.3f\ttrain sum error: %.3f\t test sum error: %.3f\n" \
"Test digits overall accuracy: %.3f\ttrain sum accuracy: %.3f\t test sum accuracy: %.3f\n" \
"Test digits overall error: %.3f\ttrain sum error: %.3f\t test sum error: %.3f\n" % \
(acc_train, acc_train1, acc_train2, error_train, error_train1, error_train2,
acc_test, acc_test1, acc_test2, error_test, error_test1, error_test2)
print(result_str)
sym_digit_network.save_result(sess, results_dir, expr, result_str)
def train(self, func, func_name='', trials=1, func_dir='results/test'):
"""Train the network to find a given function"""
x, y = generate_data(func, N_TRAIN)
# x_val, y_val = generate_data(func, N_VAL)
x_test, y_test = generate_data(func,
N_TEST,
range_min=DOMAIN_TEST[0],
range_max=DOMAIN_TEST[1])
# Setting up the symbolic regression network
x_dim = len(signature(
func).parameters) # Number of input arguments to the function
x_placeholder = tf.placeholder(shape=(None, x_dim), dtype=tf.float32)
width = len(self.activation_funcs)
n_double = functions.count_double(self.activation_funcs)
sym = SymbolicNet(self.n_layers,
funcs=self.activation_funcs,
initial_weights=[
tf.truncated_normal([x_dim, width + n_double],
stddev=init_sd_first),
tf.truncated_normal([width, width + n_double],
stddev=init_sd_middle),
tf.truncated_normal([width, width + n_double],
stddev=init_sd_middle),
tf.truncated_normal([width, 1],
stddev=init_sd_last)
])
# sym = SymbolicNet(self.n_layers, funcs=self.activation_funcs)
y_hat = sym(x_placeholder)
# Label and errors
error = tf.losses.mean_squared_error(labels=y, predictions=y_hat)
error_test = tf.losses.mean_squared_error(labels=y_test,
predictions=y_hat)
reg_loss = l12_smooth(sym.get_weights())
loss = error + self.reg_weight * reg_loss
# Set up TensorFlow graph for training
learning_rate = tf.placeholder(tf.float32)
opt = tf.train.RMSPropOptimizer(learning_rate=learning_rate)
train = opt.minimize(loss)
# Arrays to keep track of various quantities as a function of epoch
loss_list = [] # Total loss (MSE + regularization)
error_list = [] # MSE
reg_list = [] # Regularization
error_test_list = [] # Test error
error_test_final = []
eq_list = []
# Only take GPU memory as needed - allows multiple jobs on a single GPU
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
for trial in range(trials):
print("Training on function " + func_name + " Trial " +
str(trial + 1) + " out of " + str(trials))
loss_val = np.nan
# Restart training if loss goes to NaN (which happens when gradients blow up)
while np.isnan(loss_val):
sess.run(tf.global_variables_initializer())
t0 = time.time()
# First stage of training, preceded by 0th warmup stage
for i in range(self.n_epochs1 + 2000):
if i < 2000:
lr_i = self.learning_rate * 10
else:
lr_i = self.learning_rate
feed_dict = {x_placeholder: x, learning_rate: lr_i}
_ = sess.run(train, feed_dict=feed_dict)
if i % self.summary_step == 0:
loss_val, error_val, reg_val, = sess.run(
(loss, error, reg_loss), feed_dict=feed_dict)
error_test_val = sess.run(
error_test, feed_dict={x_placeholder: x_test})
print(
"Epoch: %d\tTotal training loss: %f\tTest error: %f"
% (i, loss_val, error_test_val))
loss_list.append(loss_val)
error_list.append(error_val)
reg_list.append(reg_val)
error_test_list.append(error_test_val)
if np.isnan(
loss_val
): # If loss goes to NaN, restart training
break
t1 = time.time()
# Masked network - weights below a threshold are set to 0 and frozen. This is the fine-tuning stage
sym_masked = MaskedSymbolicNet(sess, sym)
y_hat_masked = sym_masked(x_placeholder)
error_masked = tf.losses.mean_squared_error(
labels=y, predictions=y_hat_masked)
error_test_masked = tf.losses.mean_squared_error(
labels=y_test, predictions=y_hat_masked)
train_masked = opt.minimize(error_masked)
# 2nd stage of training
t2 = time.time()
for i in range(self.n_epochs2):
feed_dict = {
x_placeholder: x,
learning_rate: self.learning_rate / 10
}
_ = sess.run(train_masked, feed_dict=feed_dict)
if i % self.summary_step == 0:
loss_val, error_val = sess.run(
(loss, error_masked), feed_dict=feed_dict)
error_test_val = sess.run(
error_test_masked,
feed_dict={x_placeholder: x_test})
print(
"Epoch: %d\tTotal training loss: %f\tTest error: %f"
% (i, loss_val, error_test_val))
loss_list.append(loss_val)
error_list.append(error_val)
error_test_list.append(error_test_val)
if np.isnan(
loss_val
): # If loss goes to NaN, restart training
break
t3 = time.time()
tot_time = t1 - t0 + t3 - t2
print(tot_time)
# Print the expressions
weights = sess.run(sym_masked.get_weights())
expr = pretty_print.network(weights, self.activation_funcs,
var_names[:x_dim])
print(expr)
# Save results
trial_file = os.path.join(func_dir, 'trial%d.pickle' % trial)
results = {
"weights": weights,
"loss_list": loss_list,
"error_list": error_list,
"reg_list": reg_list,
"error_test": error_test_list,
"expr": expr,
"runtime": tot_time
}
with open(trial_file, "wb+") as f:
pickle.dump(results, f)
error_test_final.append(error_test_list[-1])
eq_list.append(expr)
return eq_list, error_test_final