def hyperparam_fitness(kernel_1, strides_1, pool_1, kernel_2, strides_2,
pool_2, kernel_3, strides_3, pool_3, kernel_4,
strides_4, pool_4, z_dimension, n_modes, n_filters_1,
n_filters_2, n_filters_3, n_filters_4, batch_size,
n_weights_fc_1, n_weights_fc_2, n_weights_fc_3):
""" Fitness function used in Gaussian Process hyperparameter optimization
Returns a value to be minimized (in this case, the total loss of the
neural network during training.
PARAMETERS:
hyperparameters to be tuned
RETURNS:
KL divergence (scalar value)
"""
# set tunable hyper-parameters
params['filter_size_r1'] = [kernel_1, kernel_2, kernel_3, kernel_4]
params['filter_size_r2'] = [kernel_1, kernel_2, kernel_3, kernel_4]
params['filter_size_q'] = [kernel_1, kernel_2, kernel_3, kernel_4]
params['n_filters_r1'] = [
n_filters_1, n_filters_2, n_filters_3, n_filters_4
]
params['n_filters_r2'] = [
n_filters_1, n_filters_2, n_filters_3, n_filters_4
]
params['n_filters_q'] = [
n_filters_1, n_filters_2, n_filters_3, n_filters_4
]
# number of filters has to be odd for some reason (this ensures that this is the case)
for filt_idx in range(len(params['n_filters_r1'])):
if (params['n_filters_r1'][filt_idx] % 3) != 0:
# keep adding 1 until filter size is divisible by 3
while (params['n_filters_r1'][filt_idx] % 3) != 0:
params['n_filters_r1'][filt_idx] += 1
params['n_filters_r2'][filt_idx] += 1
params['n_filters_q'][filt_idx] += 1
params['conv_strides_r1'] = [strides_1, strides_2, strides_3, strides_4]
params['conv_strides_r2'] = [strides_1, strides_2, strides_3, strides_4]
params['conv_strides_q'] = [strides_1, strides_2, strides_3, strides_4]
params['maxpool_r1'] = [pool_1, pool_2, pool_3, pool_4]
params['maxpool_r2'] = [pool_1, pool_2, pool_3, pool_4]
params['maxpool_q'] = [pool_1, pool_2, pool_3, pool_4]
params['pool_strides_r1'] = [pool_1, pool_2, pool_3, pool_4]
params['pool_strides_r2'] = [pool_1, pool_2, pool_3, pool_4]
params['pool_strides_q'] = [pool_1, pool_2, pool_3, pool_4]
params['z_dimension'] = z_dimension
params['n_modes'] = n_modes
params['batch_size'] = batch_size
params['n_weights_r1'] = [n_weights_fc_1, n_weights_fc_2, n_weights_fc_3]
params['n_weights_r2'] = [n_weights_fc_1, n_weights_fc_2, n_weights_fc_3]
params['n_weights_q'] = [n_weights_fc_1, n_weights_fc_2, n_weights_fc_3]
# Print the hyper-parameters.
print('kernel_1: {}'.format(kernel_1))
print('strides_1: {}'.format(strides_1))
print('pool_1: {}'.format(pool_1))
print('kernel_2: {}'.format(kernel_2))
print('strides_2: {}'.format(strides_2))
print('pool_2: {}'.format(pool_2))
print('kernel_3: {}'.format(kernel_3))
print('strides_3: {}'.format(strides_3))
print('pool_3: {}'.format(pool_3))
print('kernel_4: {}'.format(kernel_4))
print('strides_4: {}'.format(strides_4))
print('pool_4: {}'.format(pool_4))
print('z_dimension: {}'.format(z_dimension))
print('n_modes: {}'.format(n_modes))
print('n_filters_1: {}'.format(params['n_filters_r1'][0]))
print('n_filters_2: {}'.format(params['n_filters_r1'][1]))
print('n_filters_3: {}'.format(params['n_filters_r1'][2]))
print('n_filters_4: {}'.format(params['n_filters_r1'][3]))
print('batch_size: {}'.format(batch_size))
print('n_weights_r1_1: {}'.format(n_weights_fc_1))
print('n_weights_r1_2: {}'.format(n_weights_fc_2))
print('n_weights_r1_3: {}'.format(n_weights_fc_3))
print()
start_time = time.time()
print('start time: {}'.format(strftime('%X %x %Z')))
# Train model with given hyperparameters
VICI_loss, VICI_session, VICI_saver, VICI_savedir = VICI_inverse_model.train(
params, x_data_train, y_data_train, x_data_test, y_data_test,
y_data_test_noisefree, y_normscale,
"inverse_model_dir_%s/inverse_model.ckpt" % params['run_label'],
x_data_test, bounds, fixed_vals, XS_all)
end_time = time.time()
print('Run time : {} h'.format((end_time - start_time) / 3600))
# Print the loss.
print()
print("Total loss: {0:.2}".format(VICI_loss))
print()
# update variable outside of this function using global keyword
global best_loss
# save model if new best model
if VICI_loss < best_loss:
# Save model
save_path = VICI_saver.save(VICI_session, VICI_savedir)
# save hyperparameters
converged_hyperpar_dict = dict(filter_size=params['filter_size_r1'],
conv_strides=params['conv_strides_r1'],
maxpool=params['maxpool_r1'],
pool_strides=params['pool_strides_r1'],
z_dimension=params['z_dimension'],
n_modes=params['n_modes'],
n_filters=params['n_filters_r1'],
batch_size=params['batch_size'],
n_weights_fc=params['n_weights_r1'],
best_loss=best_loss)
f = open(
"inverse_model_dir_%s/converged_hyperparams.txt" %
params['run_label'], "w")
f.write(str(converged_hyperpar_dict))
f.close()
# update the best loss
best_loss = VICI_loss
# Print the loss.
print()
print("New best loss: {0:.2}".format(best_loss))
print()
# clear tensorflow session
VICI_session.close()
return VICI_loss
(params['plot_dir'], params['run_label']))
print('Did a hyperparameter search')
# otherwise, train model from user-defined hyperparameter setup
else:
# load pretrained network if wanted
if args.pretrained_loc != None:
model_loc = args.pretrained_loc
else:
model_loc = "inverse_model_dir_%s/inverse_model.ckpt" % params[
'run_label']
VICI_inverse_model.train(params, x_data_train, y_data_train,
x_data_test, y_data_test,
y_data_test_noisefree, y_normscale, model_loc,
x_data_test, bounds, fixed_vals, XS_all,
snrs_test)
# if we are now testing the network
if args.test:
# Define time series normalization scale to be using
y_normscale = 36.438613192970415
# load the testing data time series and source parameter truths
x_data_test, y_data_test_noisefree, y_data_test, _, snrs_test = load_data(
params['test_set_dir'], params['inf_pars'], load_condor=True)
# Make directory to store plots
os.system('mkdir -p %s/latest_%s' %
def train(params, x_data, y_data, siz_high_res, save_dir, plotter, y_data_test,train_files,normscales,y_data_train_noisefree,samples,pos_test,y_normscale):
x_data = x_data
y_data_train_l = y_data
# USEFUL SIZES
xsh = np.shape(x_data)
yshl1 = np.shape(y_data)[1]
ysh1 = np.shape(y_data)[1]
print(xsh,yshl1,ysh1)
#z_dimension_fm = params['z_dimensions_fw']
#n_weights_fm = params['n_weights_fw']
z_dimension = params['z_dimension']
bs = params['batch_size']
n_weights = params['n_weights']
lam = 1
graph = tf.Graph()
session = tf.Session(graph=graph)
with graph.as_default():
tf.set_random_seed(np.random.randint(0,10))
SMALL_CONSTANT = 1e-6
# PLACE HOLDERS
x_ph = tf.placeholder(dtype=tf.float32, shape=[None, xsh[1]], name="x_ph")
bs_ph = tf.placeholder(dtype=tf.int64, name="bs_ph") # batch size placeholder
yt_ph = tf.placeholder(dtype=tf.float32, shape=[None, ysh1], name="yt_ph")
# LOAD FORWARD MODEL NEURAL NETWORKS
#DEC_XYlZtoYh = OELBO_decoder_difference.VariationalAutoencoder("OELBO_decoder", ysh1, z_dimension_fm+yshl1+xsh[1], n_weights_fm) # p(Yh|X,Yl,Z)
#ENC_XYltoZ = OELBO_encoder.VariationalAutoencoder("OELBO_encoder", yshl1+xsh[1], z_dimension_fm, n_weights_fm) # p(Z|X,Yl)
#ENC_XYhYltoZ = VAE_encoder.VariationalAutoencoder("vae_encoder", xsh[1]+ysh1+yshl1, z_dimension_fm, n_weights_fm) # q(Z|X,Yl,Yh)
# LOAD VICI NEURAL NETWORKS
autoencoder = VICI_decoder.VariationalAutoencoder("VICI_decoder", xsh[1], z_dimension+ysh1, n_weights) # r(x|z,y)
autoencoder_ENC = VICI_encoder.VariationalAutoencoder("VICI_encoder", ysh1, z_dimension, n_weights) # generates params for r(z|y)
autoencoder_VAE = VICI_VAE_encoder.VariationalAutoencoder("VICI_VAE_encoder", xsh[1]+ysh1, z_dimension, n_weights) # used to sample from r(z|y)?
# DEFINE MULTI-FIDELITY FORWARD MODEL
#####################################################################################################################
SMALL_CONSTANT = 1e-6
# NORMALISE INPUTS
yl_ph_n = tf_normalise_dataset(yt_ph) # placeholder for normalised low-res y data
x_ph_n = tf_normalise_dataset(x_ph) # placeholder for normalised x data
#yl_ph_n = yt_ph
#x_ph_n = x_ph
#yl_ph_n = tf.Print(yl_ph_n, [yl_ph_n], first_n=1, summarize=10, message="Thss is yl_ph_n: ")
#x_ph_n = tf.Print(x_ph_n, [x_ph_n], first_n=1, summarize=10, message="This is x_ph_n: ")
# GET p(Z|X,Yl) - takes in x data and low res y data and returns mean and logvar of Gaussian z distribution
#zxyl_mean,zxyl_log_sig_sq = ENC_XYltoZ._calc_z_mean_and_sigma(tf.concat([x_ph_n,yl_ph_n],1))
#zxyl_mean = tf.Print(zxyl_mean, [zxyl_mean], first_n=1, summarize=10, message="Thss is zxyl_mean: ")
#zxyl_log_sig_sq = tf.Print(zxyl_log_sig_sq, [zxyl_log_sig_sq], first_n=1, summarize=10, message="Thss is zxyl_log_sig_sq: ")
# then samples z from that distribution
#rxyl_samp = ENC_XYhYltoZ._sample_from_gaussian_dist(tf.shape(x_ph_n)[0], z_dimension_fm, zxyl_mean, tf.log(tf.exp(zxyl_log_sig_sq)+SMALL_CONSTANT))
#rxyl_samp = tf.Print(rxyl_samp, [rxyl_samp], first_n=1, summarize=10, message="Thss is rxyl_samp: ")
# GET p(Yh|X,Yl,Z) FROM SAMPLES Z ~ p(Z|X,Yl)
# then decodes back to high res y data
#reconstruction_yh = DEC_XYlZtoYh.calc_reconstruction(tf.concat([x_ph_n,yl_ph_n,rxyl_samp],1))
# = tf.Print(, [], first_n=1, summarize=10, message="Thss is : ")
#reconstruction_yh = tf.Print(reconstruction_yh, [reconstruction_yh], first_n=1, summarize=10, message="Thss is reconstruction_yh: ")
#yh_diff = reconstruction_yh[0] # looks like the difference between the low res y data and the mean reconstructed high res y data
#yh_diff = tf.Print(yh_diff, [yh_diff], first_n=1, summarize=10, message="Thss is yh_diff: ")
#yh_mean = yl_ph_n+yh_diff # the mean reconstrcted high res data
#yh_mean = tf.Print(yh_mean, [yh_mean], first_n=1, summarize=10, message="Thss is yh_mean: ")
#yh_log_sig_sq = reconstruction_yh[1] # the reconstructed high res y data logvar
#yh_log_sig_sq = tf.Print(yh_log_sig_sq, [yh_log_sig_sq], first_n=1, summarize=10, message="Thss is yh_log_sig_sq: ")
# then sample something? y doesn't seem to be used anywhere after this
# this ends up being the reconstruction of the exact y data corresponding to each x data input
#y = ENC_XYhYltoZ._sample_from_gaussian_dist(tf.shape(yt_ph)[0], tf.shape(yt_ph)[1], yh_mean, yh_log_sig_sq)
#y = tf.Print(y, [y], first_n=1, summarize=10, message="Thss is y: ")
# DRAW SYNTHETIC Ys TRAINING DATA
# _, y = OM.forward_model(x_ph_amp,x_ph_ph)
##########################################################################################################################################
# GET r(z|y)
#y_ph = tf.placeholder(dtype=tf.float32, shape=[None, ysh1], name="y_ph") # placeholder for y data
#y_ph_n = tf_normalise_dataset(y_ph) # placeholder for normalised y data
#y_ph = tf.Print(y_ph, [y_ph], first_n=1, summarize=10, message="Thss is y_ph: ")
#y_ph_n = y_ph
# run inverse autoencoder to generate mean and logvar of z given y data - these are the parameters for r(z|y)
#zy_mean,zy_log_sig_sq = autoencoder_ENC._calc_z_mean_and_sigma(y_ph_n)
zy_mean,zy_log_sig_sq = autoencoder_ENC._calc_z_mean_and_sigma(yl_ph_n)
# DRAW FROM r(z|y) - given the Gaussian parameters generate z samples
rzy_samp = autoencoder_VAE._sample_from_gaussian_dist(bs_ph, z_dimension, zy_mean, zy_log_sig_sq)
# GET r(x|z,y) from r(z|y) samples
#rzy_samp_y = tf.concat([rzy_samp,y_ph_n],1)
rzy_samp_y = tf.concat([rzy_samp,yl_ph_n],1)
reconstruction_xzy = autoencoder.calc_reconstruction(rzy_samp_y)
x_mean = reconstruction_xzy[0]
x_log_sig_sq = reconstruction_xzy[1]
# KL(r(z|y)||p(z))
#latent_loss = -0.5 * tf.reduce_sum(1 + zy_log_sig_sq - tf.square(zy_mean) - tf.exp(zy_log_sig_sq), 1)
#KL = tf.reduce_mean(latent_loss)
# GET q(z|x,y)
#xy_ph = tf.concat([x_ph_n,y_ph_n],1)
xy_ph = tf.concat([x_ph_n,yl_ph_n],1)
zx_mean,zx_log_sig_sq = autoencoder_VAE._calc_z_mean_and_sigma(xy_ph)
# DRAW FROM q(z|x,y)
qzx_samp = autoencoder_VAE._sample_from_gaussian_dist(bs_ph, z_dimension, zx_mean, zx_log_sig_sq)
# GET r(x|z,y)
#qzx_samp_y = tf.concat([qzx_samp,y_ph_n],1)
qzx_samp_y = tf.concat([qzx_samp,yl_ph_n],1)
reconstruction_xzx = autoencoder.calc_reconstruction(qzx_samp_y)
x_mean_vae = reconstruction_xzx[0]
x_log_sig_sq_vae = reconstruction_xzx[1]
# COST FROM RECONSTRUCTION
normalising_factor_x_vae = - 0.5 * tf.log(SMALL_CONSTANT+tf.exp(x_log_sig_sq_vae)) - 0.5 * np.log(2 * np.pi)
square_diff_between_mu_and_x_vae = tf.square(x_mean_vae - x_ph_n)
inside_exp_x_vae = -0.5 * tf.div(square_diff_between_mu_and_x_vae,SMALL_CONSTANT+tf.exp(x_log_sig_sq_vae))
reconstr_loss_x_vae = -tf.reduce_sum(normalising_factor_x_vae + inside_exp_x_vae, 1)
cost_R_vae = tf.reduce_mean(reconstr_loss_x_vae)
# KL(q(z|x,y)||r(z|y))
v_mean = zy_mean #2
aux_mean = zx_mean #1
v_log_sig_sq = tf.log(tf.exp(zy_log_sig_sq)+SMALL_CONSTANT) #2
aux_log_sig_sq = tf.log(tf.exp(zx_log_sig_sq)+SMALL_CONSTANT) #1
v_log_sig = tf.log(tf.sqrt(tf.exp(v_log_sig_sq))) #2
aux_log_sig = tf.log(tf.sqrt(tf.exp(aux_log_sig_sq))) #1
cost_VAE_a = v_log_sig-aux_log_sig+tf.divide(tf.exp(aux_log_sig_sq)+tf.square(aux_mean-v_mean),2*tf.exp(v_log_sig_sq))-0.5
cost_VAE_b = tf.reduce_sum(cost_VAE_a,1)
KL_vae = tf.reduce_mean(cost_VAE_b) # computes the mean over all tensor elements
# THE VICI COST FUNCTION
lam_ph = tf.placeholder(dtype=tf.float32, name="lam_ph")
COST_VAE = KL_vae+cost_R_vae
COST = COST_VAE
# VARIABLES LISTS
var_list_VICI = [var for var in tf.trainable_variables() if var.name.startswith("VICI")]
var_list_ELBO = [var for var in tf.trainable_variables() if var.name.startswith("ELBO")]
# DEFINE OPTIMISER (using ADAM here)
optimizer = tf.train.AdamOptimizer(params['initial_training_rate'])
minimize = optimizer.minimize(COST,var_list = var_list_VICI)
# DRAW FROM q(x|y)
qx_samp = autoencoder_ENC._sample_from_gaussian_dist(bs_ph, xsh[1], x_mean, SMALL_CONSTANT + tf.log(tf.exp(x_log_sig_sq)))
# INITIALISE AND RUN SESSION
# init = tf.variables_initializer(var_list_VICI)
#init = tf.initialize_all_variables()
init = tf.global_variables_initializer()
session.run(init)
#saver_ELBO = tf.train.Saver(var_list_ELBO)
#saver_ELBO.restore(session,load_dir)
saver = tf.train.Saver()
KL_PLOT = np.zeros(np.int(np.round(params['num_iterations']/params['report_interval'])+1)) # vector to store test OELBO values
COST_PLOT = np.zeros(np.int(np.round(params['num_iterations']/params['report_interval'])+1)) # vector to store test VAE ELBO values
print('Training Inference Model...')
# START OPTIMISATION OF OELBO
indices_generator = batch_manager.SequentialIndexer(params['batch_size'], xsh[0])
ni = -1
test_n = 100
olvec = []
for i in range(params['num_iterations']):
next_indices = indices_generator.next_indices()
# run the session - input batchsize the x-data training batch and the y-data training batch
#yn = session.run(y, feed_dict={bs_ph:bs, x_ph:x_data[next_indices, :], yt_ph:y_data_train_l[next_indices, :]})
#session.run(minimize, feed_dict={bs_ph:bs, x_ph:x_data[next_indices, :], y_ph:yn, lam_ph:lam, yt_ph:y_data_train_l[next_indices, :]}) # minimising cost function
# Make 25 noise realizations
if params['do_extra_noise']:
x_data_train_l = x_data[next_indices,:]
y_data_train_l = y_data_train_noisefree[next_indices,:] + np.random.normal(0,1,size=(params['batch_size'],params['ndata']))
if params['do_normscale']:
y_data_train_l /= y_normscale[0]
#print('generated {} elements of new training data noise'.format(params['batch_size']))
session.run(minimize, feed_dict={bs_ph:bs, x_ph:x_data_train_l, lam_ph:lam, yt_ph:y_data_train_l}) # minimising cost function
else:
session.run(minimize, feed_dict={bs_ph:bs, x_ph:x_data[next_indices, :], lam_ph:lam, yt_ph:y_data_train_l[next_indices, :]}) # minimising cost function
if i % params['report_interval'] == 0 and i > 0:
ni = ni+1
#ynt = session.run(y, feed_dict={bs_ph:test_n, x_ph:x_data[0:test_n,:], yt_ph:y_data_train_l[0:test_n,:]})
#cost_value_vae, KL_VAE = session.run([COST_VAE, KL_vae], feed_dict={bs_ph:test_n, x_ph:x_data[0:test_n,:], y_ph:ynt, lam_ph:lam, yt_ph:y_data_train_l[0:test_n,:]})
cost_value_vae, KL_VAE = session.run([cost_R_vae, KL_vae], feed_dict={bs_ph:test_n, x_ph:x_data_train_l[0:test_n,:], lam_ph:lam, yt_ph:y_data_train_l[0:test_n,:]})
KL_PLOT[ni] = KL_VAE
COST_PLOT[ni] = cost_value_vae
# plot losses
plotter.make_loss_plot(COST_PLOT[:ni+1],KL_PLOT[:ni+1],params['report_interval'],fwd=False)
if params['print_values']==True:
print('--------------------------------------------------------------')
print('Iteration:',i)
print('Training Set -ELBO:',cost_value_vae)
print('KL Divergence:',KL_VAE)
if i % params['plot_interval'] == 0 and i>0:
# The trained inverse model weights can then be used to infer a probability density of solutions given new measurements
_, _, XS, _, _ = VICI_inverse_model.run(params, y_data_test, np.shape(x_data)[1], "inverse_model_dir_%s/inverse_model.ckpt" % params['run_label'])
# Convert XS back to unnormalized version
if params['do_normscale']:
for m in range(params['ndim_x']):
XS[:,m,:] = XS[:,m,:]*normscales[m]
# Generate final results plots
plotter = plots.make_plots(params,samples,XS,pos_test)
# Make corner plots
plotter.make_corner_plot(sampler='dynesty1')
# Make KL plot
# plotter.gen_kl_plots(VICI_inverse_model,y_data_test,x_data,normscales)
# Make pp plot
# plotter.plot_pp(VICI_inverse_model,y_data_train_l,x_data_train,0,normscales)
if i % params['save_interval'] == 0 and i > 0:
# Save model
save_path = saver.save(session,save_dir)
return COST_PLOT, KL_PLOT, train_files
normscales = [normscales[0],normscales[2],normscales[3],normscales[4]]#,normscales[5]]
x_data_train, y_data_train_l, y_data_train_h = x_data_train_h, y_data_train_lh, y_data_train_lh
# Make directory for plots
plots.make_dirs(params,params['plot_dir'][0])
# Declare plot class variables
plotter = plots.make_plots(params,samples,None,pos_test)
# Plot test sample time series
plotter.plot_testdata((y_data_test_h),sig_test,params['r']**2,params['plot_dir'])
# Train model
if not params['do_only_test']:
loss_inv, kl_inv, train_files = VICI_inverse_model.train(params, x_data_train, y_data_train_l, np.shape(y_data_train_h)[1], "inverse_model_dir_%s/inverse_model.ckpt" % params['run_label'], plotter, y_data_test_h,train_files,normscales,y_data_train_noisefree,samples,pos_test,y_normscale) # This trains the inverse model to recover posteriors using the forward model weights stored in forward_model_dir/forward_model.ckpt and saves the inverse model weights in inverse_model_dir/inverse_model.ckpt
# Test model
if params['do_only_test']:
# The trained inverse model weights can then be used to infer a probability density of solutions given new measurements
xm, xsx, XS, pmax, _ = VICI_inverse_model.run(params, y_data_test_h, np.shape(x_data_train)[1], "inverse_model_dir_%s/inverse_model.ckpt" % params['run_label']) # This runs the trained model using the weights stored in inverse_model_dir/inverse_model.ckpt
# TODO: remove this later
# normscales = [79.99998770493086, 0.5999996662139893, 79.95626661877512, 2999.9996467099486]
# Convert XS back to unnormalized version
if params['do_normscale']:
for m in range(params['ndim_x']):
XS[:,m,:] = XS[:,m,:] * normscales[m]
# Generate final results plots
def resume_training(params, x_data, y_data_l, siz_high_res, save_dir, train_files,normscales,y_data_train_noisefree,y_normscale):
x_data = x_data
y_data_train_l = y_data_l
# USEFUL SIZES
xsh = np.shape(x_data)
yshl1 = np.shape(y_data_l)[1]
ysh1 = siz_high_res
#z_dimension_fm = params['z_dimensions_fw']
#n_weights_fm = params['n_weights_fw']
z_dimension = params['z_dimension']
bs = params['batch_size']
n_weights = params['n_weights']
lam = 1
# Allow GPU growth
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
graph = tf.Graph()
session = tf.Session(graph=graph,config=config)
with graph.as_default():
tf.set_random_seed(np.random.randint(0,10))
SMALL_CONSTANT = 1e-6
# PLACE HOLDERS
x_ph = tf.placeholder(dtype=tf.float32, shape=[None, xsh[1]], name="x_ph")
bs_ph = tf.placeholder(dtype=tf.int64, name="bs_ph") # batch size placeholder
yt_ph = tf.placeholder(dtype=tf.float32, shape=[None, ysh1], name="yt_ph")
# LOAD FORWARD MODEL NEURAL NETWORKS
#DEC_XYlZtoYh = OELBO_decoder_difference.VariationalAutoencoder("OELBO_decoder", ysh1, z_dimension_fm+yshl1+xsh[1], n_weights_fm) # p(Yh|X,Yl,Z)
#ENC_XYltoZ = OELBO_encoder.VariationalAutoencoder("OELBO_encoder", yshl1+xsh[1], z_dimension_fm, n_weights_fm) # p(Z|X,Yl)
#ENC_XYhYltoZ = VAE_encoder.VariationalAutoencoder("vae_encoder", xsh[1]+ysh1+yshl1, z_dimension_fm, n_weights_fm) # q(Z|X,Yl,Yh)
# LOAD VICI NEURAL NETWORKS
autoencoder = VICI_decoder.VariationalAutoencoder("VICI_decoder", xsh[1], z_dimension+ysh1, n_weights)
autoencoder_ENC = VICI_encoder.VariationalAutoencoder("VICI_encoder", ysh1, z_dimension, n_weights)
autoencoder_VAE = VICI_VAE_encoder.VariationalAutoencoder("VICI_VAE_encoder", xsh[1]+ysh1, z_dimension, n_weights)
# DEFINE MULTI-FIDELITY FORWARD MODEL
#####################################################################################################################
SMALL_CONSTANT = 1e-6
# NORMALISE INPUTS
yl_ph_n = tf_normalise_dataset(yt_ph)
x_ph_n = tf_normalise_dataset(x_ph)
#yl_ph_n = yt_ph
#x_ph_n = x_ph
# GET p(Z|X,Yl)
#zxyl_mean,zxyl_log_sig_sq = ENC_XYltoZ._calc_z_mean_and_sigma(tf.concat([x_ph_n,yl_ph_n],1))
#rxyl_samp = ENC_XYhYltoZ._sample_from_gaussian_dist(tf.shape(x_ph_n)[0], z_dimension_fm, zxyl_mean, tf.log(tf.exp(zxyl_log_sig_sq)+SMALL_CONSTANT))
# GET p(Yh|X,Yl,Z) FROM SAMPLES Z ~ p(Z|X,Yl)
#reconstruction_yh = DEC_XYlZtoYh.calc_reconstruction(tf.concat([x_ph_n,yl_ph_n,rxyl_samp],1))
#yh_diff = reconstruction_yh[0]
#yh_mean = yl_ph_n+yh_diff
#yh_log_sig_sq = reconstruction_yh[1]
#y = ENC_XYhYltoZ._sample_from_gaussian_dist(tf.shape(yt_ph)[0], tf.shape(yt_ph)[1], yh_mean, yh_log_sig_sq)
# DRAW SYNTHETIC Ys TRAINING DATA
# _, y = OM.forward_model(x_ph_amp,x_ph_ph)
##########################################################################################################################################
# GET r(z|y)
#y_ph = tf.placeholder(dtype=tf.float32, shape=[None, ysh1], name="y_ph")
#y_ph_n = tf_normalise_dataset(y_ph)
#y_ph_n = y_ph
zy_mean,zy_log_sig_sq = autoencoder_ENC._calc_z_mean_and_sigma(yl_ph_n)
# DRAW FROM r(z|y)
rzy_samp = autoencoder_VAE._sample_from_gaussian_dist(bs_ph, z_dimension, zy_mean, zy_log_sig_sq)
# GET r(x|z,y) from r(z|y) samples
rzy_samp_y = tf.concat([rzy_samp,yl_ph_n],1)
reconstruction_xzy = autoencoder.calc_reconstruction(rzy_samp_y)
x_mean = reconstruction_xzy[0]
x_log_sig_sq = reconstruction_xzy[1]
# KL(r(z|y)||p(z))
latent_loss = -0.5 * tf.reduce_sum(1 + zy_log_sig_sq - tf.square(zy_mean) - tf.exp(zy_log_sig_sq), 1)
KL = tf.reduce_mean(latent_loss)
# GET q(z|x,y)
xy_ph = tf.concat([x_ph_n,yl_ph_n],1)
zx_mean,zx_log_sig_sq = autoencoder_VAE._calc_z_mean_and_sigma(xy_ph)
# DRAW FROM q(z|x,y)
qzx_samp = autoencoder_VAE._sample_from_gaussian_dist(bs_ph, z_dimension, zx_mean, zx_log_sig_sq)
# GET r(x|z,y)
qzx_samp_y = tf.concat([qzx_samp,yl_ph_n],1)
reconstruction_xzx = autoencoder.calc_reconstruction(qzx_samp_y)
x_mean_vae = reconstruction_xzx[0]
x_log_sig_sq_vae = reconstruction_xzx[1]
# COST FROM RECONSTRUCTION
normalising_factor_x_vae = - 0.5 * tf.log(SMALL_CONSTANT+tf.exp(x_log_sig_sq_vae)) - 0.5 * np.log(2 * np.pi)
square_diff_between_mu_and_x_vae = tf.square(x_mean_vae - x_ph_n)
inside_exp_x_vae = -0.5 * tf.div(square_diff_between_mu_and_x_vae,SMALL_CONSTANT+tf.exp(x_log_sig_sq_vae))
reconstr_loss_x_vae = -tf.reduce_sum(normalising_factor_x_vae + inside_exp_x_vae, 1) # Take sum along axis 1
cost_R_vae = tf.reduce_mean(reconstr_loss_x_vae) # Take mean along the diagonal
# KL(q(z|x,y)||r(z|y))
v_mean = zy_mean #2
aux_mean = zx_mean #1
v_log_sig_sq = tf.log(tf.exp(zy_log_sig_sq)+SMALL_CONSTANT) #2
aux_log_sig_sq = tf.log(tf.exp(zx_log_sig_sq)+SMALL_CONSTANT) #1
v_log_sig = tf.log(tf.sqrt(tf.exp(v_log_sig_sq))) #2
aux_log_sig = tf.log(tf.sqrt(tf.exp(aux_log_sig_sq))) #1
cost_VAE_a = v_log_sig-aux_log_sig+tf.divide(tf.exp(aux_log_sig_sq)+tf.square(aux_mean-v_mean),2*tf.exp(v_log_sig_sq))-0.5
cost_VAE_b = tf.reduce_sum(cost_VAE_a,1)
KL_vae = tf.reduce_mean(cost_VAE_b)
# THE VICI COST FUNCTION
lam_ph = tf.placeholder(dtype=tf.float32, name="lam_ph")
COST_VAE = KL_vae+cost_R_vae
COST = COST_VAE
# VARIABLES LISTS
var_list_VICI = [var for var in tf.trainable_variables() if var.name.startswith("VICI")]
var_list_ELBO = [var for var in tf.trainable_variables() if var.name.startswith("ELBO")]
# DEFINE OPTIMISER (using ADAM here)
optimizer = tf.train.AdamOptimizer(params['initial_training_rate'])
minimize = optimizer.minimize(COST,var_list = var_list_VICI)
# DRAW FROM q(x|y)
qx_samp = autoencoder_ENC._sample_from_gaussian_dist(bs_ph, xsh[1], x_mean, SMALL_CONSTANT + tf.log(tf.exp(x_log_sig_sq)))
# INITIALISE AND RUN SESSION
# init = tf.variables_initializer(var_list_VICI)
init = tf.initialize_all_variables()
session.run(init)
#saver_ELBO = tf.train.Saver(var_list_ELBO)
#saver_ELBO.restore(session,load_dir)
saver = tf.train.Saver(var_list_VICI)
saver.restore(session,save_dir)
KL_PLOT = np.zeros(np.int(np.round(params['num_iterations']/params['report_interval'])+1)) # vector to store test OELBO values
COST_PLOT = np.zeros(np.int(np.round(params['num_iterations']/params['report_interval'])+1)) # vector to store test VAE ELBO values
print('Training Inference Model...')
# START OPTIMISATION OF OELBO
indices_generator = batch_manager.SequentialIndexer(params['batch_size'], xsh[0])
ni = -1
test_n = 100
for i in range(params['num_iterations']):
next_indices = indices_generator.next_indices()
# Make 25 noise realizations
if params['do_extra_noise']:
x_data_train_l = x_data[next_indices,:]
y_data_train_l = y_data_train_noisefree[next_indices,:] + np.random.normal(0,1,size=(params['batch_size'],params['ndata']))
y_data_train_l /= y_normscale[0]
#print('generated {} elements of new training data noise'.format(params['batch_size']))
session.run(minimize, feed_dict={bs_ph:bs, x_ph:x_data_train_l, lam_ph:lam, yt_ph:y_data_train_l}) # minimising cost function
else:
#yn = session.run(y, feed_dict={bs_ph:bs, x_ph:x_data[next_indices, :], yt_ph:y_data_train_l[next_indices, :]})
#session.run(minimize, feed_dict={bs_ph:bs, x_ph:x_data[next_indices, :], y_ph:yn, lam_ph:lam, yt_ph:y_data_train_l[next_indices, :]}) # minimising cost function
session.run(minimize, feed_dict={bs_ph:bs, x_ph:x_data[next_indices, :], lam_ph:lam, yt_ph:y_data_train_l[next_indices, :]}) # minimising cost function
if i % params['report_interval'] == 0:
ni = ni+1
#ynt = session.run(y, feed_dict={bs_ph:test_n, x_ph:x_data[0:test_n,:], yt_ph:y_data_train_l[0:test_n,:]})
#cost_value_vae, KL_VAE = session.run([COST_VAE, KL_vae], feed_dict={bs_ph:test_n, x_ph:x_data[0:test_n,:], y_ph:ynt, lam_ph:lam, yt_ph:y_data_train_l[0:test_n,:]})
cost_value_vae, KL_VAE = session.run([COST_VAE, KL_vae], feed_dict={bs_ph:test_n, x_ph:x_data[0:test_n,:], lam_ph:lam, yt_ph:y_data_train_l[0:test_n,:]})
KL_PLOT[ni] = KL_VAE
COST_PLOT[ni] = -cost_value_vae
# plot losses
plotter.make_loss_plot(COST_PLOT[:ni+1],KL_PLOT[:ni+1],params['report_interval'],fwd=False)
if params['print_values']==True:
print('--------------------------------------------------------------')
print('Iteration:',i)
print('Training Set ELBO:',-cost_value_vae)
print('KL Divergence:',KL_VAE)
if i % params['plot_interval'] == 0 and i>0:
# The trained inverse model weights can then be used to infer a probability density of solutions given new measurements
_, _, XS, _ = VICI_inverse_model.run(params, y_data_test_h, np.shape(x_data_train)[1], "inverse_model_dir_%s/inverse_model.ckpt" % params['run_label'])
# Convert XS back to unnormalized version
if params['do_normscale']:
for m in range(params['ndim_x']):
XS[:,m,:] = XS[:,m,:]*normscales[m]
# Make KL plot
plotter.gen_kl_plots(VICI_inverse_model,y_data_test_h,x_data_train,normscales)
# Make corner plots
plotter.make_corner_plot(sampler='dynesty1')
# Make KL plot
plotter.gen_kl_plots(VICI_inverse_model,y_data_test_h,x_data_train,normscales)
# Make pp plot
# plotter.plot_pp(VICI_inverse_model,y_data_train_l,x_data_train,0,normscales)
if i % params['save_interval'] == 0:
save_path = saver.save(session,save_dir)
return COST_PLOT, KL_PLOT, train_files