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