def main(_):
util.print('Loading model {} with dataset {}.'.format(
FLAGS.model, FLAGS.dataset))
if FLAGS.model == 'radon':
model_config = models.get_radon(state_code=FLAGS.dataset)
elif FLAGS.model == 'radon_stddvs':
model_config = models.get_radon_model_stddvs(state_code=FLAGS.dataset)
elif FLAGS.model == '8schools':
model_config = models.get_eight_schools()
elif FLAGS.model == 'german_credit_gammascale':
model_config = models.get_german_credit_gammascale()
elif FLAGS.model == 'german_credit_lognormalcentered':
model_config = models.get_german_credit_lognormalcentered()
else:
raise Exception('unknown model {}'.format(FLAGS.model))
description = FLAGS.model + '_{}'.format(FLAGS.dataset)
experiments_dir = os.path.join(
FLAGS.results_dir,
'num_leapfrog_steps={}'.format(FLAGS.num_leapfrog_steps))
if not tf.gfile.Exists(experiments_dir):
tf.gfile.MakeDirs(experiments_dir)
if FLAGS.method == 'baseline':
run_baseline(description,
model_config=model_config,
experiments_dir=experiments_dir,
num_samples=FLAGS.num_samples,
burnin=FLAGS.burnin,
num_adaptation_steps=FLAGS.num_adaptation_steps,
num_optimization_steps=FLAGS.num_optimization_steps,
tau=FLAGS.tau,
num_leapfrog_steps=FLAGS.num_leapfrog_steps,
description=description)
elif FLAGS.method == 'vip':
run_vip(description,
model_config=model_config,
experiments_dir=experiments_dir,
use_iaf_posterior=FLAGS.use_iaf_posterior,
num_samples=FLAGS.num_samples,
burnin=FLAGS.burnin,
num_adaptation_steps=FLAGS.num_adaptation_steps,
num_optimization_steps=FLAGS.num_optimization_steps,
num_mc_samples=FLAGS.num_mc_samples,
tau=FLAGS.tau,
num_leapfrog_steps=FLAGS.num_leapfrog_steps,
description=description)
else:
raise Exception('No such method')
def run_baseline(experiment_name,
model_config,
experiments_dir,
description,
num_samples=2000,
burnin=1000,
num_adaptation_steps=500,
num_leapfrog_steps=4,
num_optimization_steps=2000,
tau=1.):
folder_path = os.path.join(experiments_dir, experiment_name)
if not tf.gfile.Exists(folder_path):
tf.gfile.MakeDirs(folder_path)
spec = 'Model name: {}\n{}\n\n'.format(experiment_name, description)
spec += '{} samples\n'.format(num_samples)
spec += '{} burnin\n'.format(burnin)
spec += '{} adaptation steps\n'.format(num_adaptation_steps)
spec += '{} leapfrog steps\n'.format(num_leapfrog_steps)
spec += '{} optimization steps\n'.format(num_optimization_steps)
if tau != 1.0:
spec += 'tau = {}\n'.format(tau)
util.print('\nRunning HMC-CP...')
util.print(spec)
results = algs.run_centered_hmc(
model_config=model_config,
num_samples=num_samples,
burnin=burnin,
num_leapfrog_steps=num_leapfrog_steps,
num_adaptation_steps=num_adaptation_steps,
num_optimization_steps=num_optimization_steps)
util.print(' ess / leapfrogs: {}'.format(
min([np.amin(ess) for ess in results['ess']]) / num_leapfrog_steps))
util.print(' avg ess / leapfrogs: {}'.format(
np.mean([np.mean(ess)
for ess in results['ess']]) / num_leapfrog_steps))
util.print(' acceptance rate: {}'.format(results['acceptance_rate']))
util.print(' VI ran for: {} minutes'.format(results['vi_time'] / 60.))
util.print(' HMC ran for: {} minutes'.format(results['sampling_time'] /
60.))
step_size_cp = results['step_size']
results['algorithm'] = 'HMC-CP'
results['num_leapfrog_steps'] = num_leapfrog_steps
save_path = os.path.join(folder_path, results['algorithm'])
pickle_results(results, save_path)
util.print('\nRunning HMC-NCP...')
results = algs.run_noncentered_hmc(
model_config=model_config,
num_samples=num_samples,
burnin=burnin,
num_leapfrog_steps=num_leapfrog_steps,
num_adaptation_steps=num_adaptation_steps,
num_optimization_steps=num_optimization_steps)
util.print(' ess / leapfrogs: {}'.format(
min([np.amin(ess) for ess in results['ess']]) / num_leapfrog_steps))
util.print(' avg ess / leapfrogs: {}'.format(
np.mean([np.mean(ess)
for ess in results['ess']]) / num_leapfrog_steps))
util.print(' acceptance rate: {}'.format(results['acceptance_rate']))
util.print(' VI ran for: {} minutes'.format(results['vi_time'] / 60.))
util.print(' HMC ran for: {} minutes'.format(results['sampling_time'] /
60.))
step_size_ncp = results['step_size']
results['algorithm'] = 'HMC-NCP'
results['num_leapfrog_steps'] = num_leapfrog_steps
save_path = os.path.join(folder_path, results['algorithm'])
pickle_results(results, save_path)
util.print('\nRunning iHMC...')
results = algs.run_interleaved_hmc(model_config=model_config,
num_samples=num_samples,
burnin=burnin,
num_leapfrog_steps=num_leapfrog_steps,
step_size_cp=step_size_cp,
step_size_ncp=step_size_ncp)
util.print(' ess / leapfrogs: {}'.format(
min([np.amin(ess) for ess in results['ess']]) / num_leapfrog_steps))
util.print(' avg ess / leapfrogs: {}'.format(
np.mean([np.mean(ess)
for ess in results['ess']]) / num_leapfrog_steps))
util.print(' acceptance rate: {}'.format(results['acceptance_rate']))
util.print(' HMC ran for: {} minutes'.format(results['sampling_time'] /
60.))
results['algorithm'] = 'iHMC'
results['num_leapfrog_steps'] = num_leapfrog_steps
save_path = os.path.join(folder_path, results['algorithm'])
pickle_results(results, save_path)
def run_vip(experiment_name,
model_config,
experiments_dir,
num_samples=2000,
burnin=1000,
num_adaptation_steps=500,
num_leapfrog_steps=4,
num_optimization_steps=2000,
num_mc_samples=32,
tau=1.,
use_iaf_posterior=False,
description=''):
folder_path = os.path.join(experiments_dir, experiment_name)
if not tf.gfile.Exists(folder_path):
tf.gfile.MakeDirs(folder_path)
spec = 'Model name: {}\n{}\n\n'.format(experiment_name, description)
spec += '{} samples\n'.format(num_samples)
spec += '{} burnin\n'.format(burnin)
spec += '{} adaptation steps\n'.format(num_adaptation_steps)
spec += '{} leapfrog steps\n'.format(num_leapfrog_steps)
spec += '{} optimization steps\n'.format(num_optimization_steps)
spec += '{} mc samples\n'.format(num_mc_samples)
if tau != 1.0:
spec += 'tau = {}\n'.format(tau)
util.print('\nRunning c-VIP-HMC...')
results = algs.run_vip_hmc_continuous(
model_config=model_config,
num_samples=num_samples,
burnin=burnin,
use_iaf_posterior=use_iaf_posterior,
num_leapfrog_steps=num_leapfrog_steps,
num_adaptation_steps=num_adaptation_steps,
num_optimization_steps=num_optimization_steps,
num_mc_samples=num_mc_samples,
experiments_dir=experiments_dir,
tau=tau,
description=description + 'c-VIP')
if not use_iaf_posterior:
util.print(' ess / leapfrogs: {}'.format(
min([np.amin(ess)
for ess in results['ess']]) / num_leapfrog_steps))
util.print(' avg ess / leapfrogs: {}'.format(
np.mean([np.mean(ess)
for ess in results['ess']]) / num_leapfrog_steps))
util.print(' acceptance rate: {}'.format(results['acceptance_rate']))
util.print(' VI ran for: {} minutes'.format(results['vi_time'] /
60.))
util.print(' HMC ran for: {} minutes'.format(
results['sampling_time'] / 60.))
results['algorithm'] = 'c-VIP-HMC'
results['num_leapfrog_steps'] = num_leapfrog_steps
save_path = os.path.join(folder_path, results['algorithm'])
pickle_results(results, save_path)
parameterisation = results['parameterisation']
discrete_parameterisation = collections.OrderedDict([
(key, (np.array(parameterisation[key]) >= 0.5).astype(np.float32))
for key in parameterisation.keys()
])
vip_to_centered = model_config.make_to_centered(
**discrete_parameterisation)
model_config_dvip = model_config._replace(to_centered=vip_to_centered)
util.print('\nRunning d-VIP-HMC...')
results = algs.run_vip_hmc_discrete(
model_config=model_config_dvip,
parameterisation=discrete_parameterisation,
num_samples=num_samples,
burnin=burnin,
num_leapfrog_steps=num_leapfrog_steps,
num_adaptation_steps=num_adaptation_steps,
num_optimization_steps=num_optimization_steps)
util.print(' ess / leapfrogs: {}'.format(
min([np.amin(ess) for ess in results['ess']]) / num_leapfrog_steps))
util.print(' avg ess / leapfrogs: {}'.format(
np.mean([np.mean(ess)
for ess in results['ess']]) / num_leapfrog_steps))
util.print(' acceptance rate: {}'.format(results['acceptance_rate']))
util.print(' HMC ran for: {} minutes'.format(results['sampling_time'] /
60.))
results['algorithm'] = 'd-VIP-HMC'
results['num_leapfrog_steps'] = num_leapfrog_steps
save_path = os.path.join(folder_path, results['algorithm'])
pickle_results(results, save_path)
def run_vip_hmc_continuous(model_config,
num_samples=2000,
burnin=1000,
use_iaf_posterior=False,
num_leapfrog_steps=4,
num_adaptation_steps=500,
num_optimization_steps=2000,
num_mc_samples=32,
tau=1.,
do_sample=True,
description='',
experiments_dir=''):
tf.reset_default_graph()
if use_iaf_posterior:
# IAF posterior doesn't give us stddevs for step sizes for HMC (we could
# extract them by sampling but I haven't implemented that), and we mostly
# care about it for ELBOs anyway.
do_sample = False
init_val_loc = tf.placeholder('float', shape=())
init_val_scale = tf.placeholder('float', shape=())
(learnable_parameters,
learnable_parametrisation, _) = ed_transforms.make_learnable_parametrisation(
init_val_loc=init_val_loc, init_val_scale=init_val_scale, tau=tau)
def model_vip(*params):
with ed.interception(learnable_parametrisation):
return model_config.model(*params)
log_joint_vip = ed.make_log_joint_fn(model_vip)
with ed.tape() as model_tape:
_ = model_vip(*model_config.model_args)
param_shapes = collections.OrderedDict()
target_vip_kwargs = {}
for param in model_tape.keys():
if param not in model_config.observed_data.keys():
param_shapes[param] = model_tape[param].shape
else:
target_vip_kwargs[param] = model_config.observed_data[param]
def target_vip(*param_args):
i = 0
for param in model_tape.keys():
if param not in model_config.observed_data.keys():
target_vip_kwargs[param] = param_args[i]
i = i + 1
return log_joint_vip(*model_config.model_args, **target_vip_kwargs)
full_kwargs = collections.OrderedDict(model_config.observed_data.items())
full_kwargs['parameterisation'] = collections.OrderedDict()
for k in learnable_parameters.keys():
full_kwargs['parameterisation'][k] = learnable_parameters[k]
if use_iaf_posterior:
elbo = util.get_iaf_elbo(
target_vip,
num_mc_samples=num_mc_samples,
param_shapes=param_shapes)
variational_parameters = {}
else:
elbo, variational_parameters = util.get_mean_field_elbo(
model_vip,
target_vip,
num_mc_samples=num_mc_samples,
model_args=model_config.model_args,
vi_kwargs=full_kwargs)
vip_step_size_approx = util.get_approximate_step_size(
variational_parameters, num_leapfrog_steps)
##############################################################################
best_elbo = None
model_dir = os.path.join(experiments_dir,
str(description + '_' + model_config.model.__name__))
if not tf.gfile.Exists(model_dir):
tf.gfile.MakeDirs(model_dir)
saver = tf.train.Saver()
dir_save = os.path.join(model_dir, 'saved_params_{}'.format(gen_id()))
if not tf.gfile.Exists(dir_save):
tf.gfile.MakeDirs(dir_save)
best_lr = None
best_init_loc = None
best_init_scale = None
learning_rate_ph = tf.placeholder(shape=[], dtype=tf.float32)
learning_rate = tf.Variable(learning_rate_ph, trainable=False)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
train = optimizer.minimize(-elbo)
init = tf.global_variables_initializer()
learning_rates = [0.003, 0.01, 0.01, 0.1, 0.003, 0.01]
if use_iaf_posterior:
learning_rates = [3e-5, 1e-4, 3e-4, 1e-4]
start_time = time.time()
for learning_rate_val in learning_rates:
for init_loc in [0.]: #, 10., -10.]:
for init_scale in [init_loc]:
timeline = []
with tf.Session() as sess:
init.run(feed_dict={init_val_loc: init_loc,
init_val_scale: init_scale,
learning_rate_ph: learning_rate_val})
this_timeline = []
for i in range(num_optimization_steps):
_, e = sess.run([train, elbo])
if np.isnan(e):
util.print('got NaN in ELBO optimization, stopping...')
break
this_timeline.append(e)
this_elbo = np.mean(this_timeline[-100:])
info_str = ('finished cVIP optimization with elbo {} vs '
'best ELBO {}'.format(this_elbo, best_elbo))
util.print(info_str)
if best_elbo is None or best_elbo < this_elbo:
best_elbo = this_elbo
timeline = this_timeline
vals = sess.run(list(learnable_parameters.values()))
learned_reparam = collections.OrderedDict(
zip(learnable_parameters.keys(), vals))
vals = sess.run(list(variational_parameters.values()))
learned_variational_params = collections.OrderedDict(
zip(variational_parameters.keys(), vals))
util.print('learned params {}'.format(learned_reparam))
util.print('learned variational params {}'.format(
learned_variational_params))
_ = saver.save(sess, dir_save)
best_lr = learning_rate
best_init_loc = init_loc
best_init_scale = init_scale
vi_time = time.time() - start_time
util.print('BEST: LR={}, init={}, {}'.format(best_lr, best_init_loc,
best_init_scale))
util.print('ELBO: {}'.format(best_elbo))
to_centered = model_config.make_to_centered(**learned_reparam)
results = collections.OrderedDict()
results['elbo'] = best_elbo
with tf.Session() as sess:
saver.restore(sess, dir_save)
results['vp'] = learned_variational_params
if do_sample:
vip_step_size_init = sess.run(vip_step_size_approx)
vip_step_size = [tf.get_variable(
name='step_size_vip'+str(i),
initializer=np.array(vip_step_size_init[i], dtype=np.float32),
use_resource=True, # For TFE compatibility.
trainable=False) for i in range(len(vip_step_size_init))]
kernel_vip = mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=target_vip,
step_size=vip_step_size,
num_leapfrog_steps=num_leapfrog_steps,
step_size_update_fn=mcmc.make_simple_step_size_update_policy(
num_adaptation_steps=num_adaptation_steps, target_rate=0.85))
states, kernel_results_vip = mcmc.sample_chain(
num_results=num_samples,
num_burnin_steps=burnin,
current_state=[
tf.zeros(param_shapes[param]) for param in param_shapes.keys()
],
kernel=kernel_vip,
num_steps_between_results=1)
states_vip = transform_mcmc_states(states, to_centered)
init_again = tf.global_variables_initializer()
init_again.run(feed_dict={
init_val_loc: best_init_loc, init_val_scale: best_init_scale,
learning_rate_ph: 1.0}) # learning rate doesn't matter for HMC.
ess_vip = tfp.mcmc.effective_sample_size(states_vip)
start_time = time.time()
samples, is_accepted, ess, ss_vip, log_accept_ratio = sess.run(
(states_vip, kernel_results_vip.is_accepted, ess_vip,
kernel_results_vip.extra.step_size_assign,
kernel_results_vip.log_accept_ratio))
sampling_time = time.time() - start_time
results['samples'] = collections.OrderedDict()
results['is_accepted'] = is_accepted
results['acceptance_rate'] = np.sum(is_accepted) * 100. / float(
num_samples)
results['ess'] = ess
results['sampling_time'] = sampling_time
results['log_accept_ratio'] = log_accept_ratio
results['step_size'] = [s[0] for s in ss_vip]
i = 0
for param in param_shapes.keys():
results['samples'][param] = samples[i]
i = i + 1
# end if
results['parameterisation'] = collections.OrderedDict()
i = 0
for param in param_shapes.keys():
name_a = param[:-5] + 'a'
name_b = param[:-5] + 'b'
try:
results['parameterisation'][name_a] = learned_reparam[name_a]
results['parameterisation'][name_b] = learned_reparam[name_b]
except KeyError:
continue
i = i + 1
results['elbo_timeline'] = timeline
results['vi_time'] = vi_time
results['init_pos'] = best_init_loc
return results