def evaluate_metrics(samples, sample_time, suffix): samples = samples[self.args['burn_in']:] np.save( os.path.join(self.logs['results'], 'samples_%s.npy' % suffix), samples) self._save_logs('results.txt', 'sample time %s' % suffix, sample_time) print('evaluating samples') bmess = batch_means_ess(samples) min_bmess = np.mean(np.min(bmess, axis=1), axis=0) std_bmess = np.std(np.min(bmess, axis=1), axis=0) gr = gelman_rubin_diagnostic(samples) acc_rate = acceptance_rate_2(samples) self._save_logs('results.txt', 'num_samples_%s:' % suffix, samples.shape[0]) self._save_logs('results.txt', 'ess_%s' % suffix, bmess) self._save_logs('results.txt', 'min_ess_%s' % suffix, min_bmess) self._save_logs('results.txt', 'std_ess_%s' % suffix, std_bmess) self._save_logs('results.txt', 'gelman_rubin_%s:' % suffix, gr) self._save_logs('results.txt', 'acceptance_rate_%s:' % suffix, acc_rate) self._trace_plot(samples) eval, eval_std = self._evaluate(samples) self._save_logs('results.txt', 'eval_metric_%s' % suffix, eval) self._save_logs('results.txt', 'eval_metric_std_%s' % suffix, eval_std)
def run(self): print('starting training') losses, train_time = self.sampler.train(**self.args) np.save(os.path.join(self.logs['results'], 'losses.npy'), np.array(losses)) self._plot(losses[1:], 'losses.png') self._save_logs('results.txt', 'train time', train_time) print('drawing samples') samples, sample_time = self.sampler.sample(**self.args) samples = samples[self.args['burn_in']:] np.save(os.path.join(self.logs['results'], 'samples.npy'), samples) self._save_logs('results.txt', 'sample time', sample_time) print('evaluating samples') ess = batch_means_ess(samples) gr = gelman_rubin_diagnostic(samples) acc_rate = acceptance_rate_2(samples) np.save(os.path.join(self.logs['results'], 'ess.npy'), ess) min_ess = np.mean(np.min(ess, axis=1), axis=0) std_ess = np.std(np.min(ess, axis=1), axis=0) self._save_logs('results.txt', 'min_ess', min_ess) self._save_logs('results.txt', 'std_ess', std_ess) self._save_logs('results.txt', 'num_samples:', samples.size) self._save_logs('results.txt', 'gelman_rubin:', gr) self._save_logs('results.txt', 'acceptance_rate:', acc_rate) self._trace_plot(samples) eval, eval_std = self._evaluate(samples) self._save_logs('results.txt', 'eval_metric', eval) self._save_logs('results.txt', 'eval_metric_std', eval_std) if samples.shape[2] == 2: self._plot2d(samples)
def train(self, sess=None, d_iters=5, epoch_size=1000, max_iters=100000, bootstrap_steps=5000, bootstrap_burn_in=1000, bootstrap_batch_size=32, bootstrap_discard_ratio=0.5, evaluate_steps=10000, evaluate_burn_in=5000, evaluate_batch_size=32, nice_steps=1, hmc_epochs=5, batch_size=32, use_hmc=False, save_path=None, log_freq=100, plot=True, save_freq=2000, hmc_steps=100, hmc_step_size=0.001, **kwargs): """ Train the NICE proposal using adversarial training. :param d_iters: number of discrtiminator iterations for each generator iteration :param epoch_size: how many iteration for each bootstrap step :param log_freq: how many iterations for each log on screen :param max_iters: max number of iterations for training :param bootstrap_steps: how many steps for each bootstrap :param bootstrap_burn_in: how many burn in steps for each bootstrap :param bootstrap_batch_size: # of chains for each bootstrap :param bootstrap_discard_ratio: ratio for discarding previous samples :param evaluate_steps: how many steps to evaluate performance :param evaluate_burn_in: how many burn in steps to evaluate performance :param evaluate_batch_size: # of chains for evaluating performance :param nice_steps: Experimental. num of steps for running the nice proposal before MH. For now do not use larger than 1. :param hmc_epochs: how many bootstrap epochs to use HMC for before switiching to the model :return: """ sess.run(self.init_op) saver = tf.train.Saver(max_to_keep=1) def _feed_dict(bs): return { self.z: self.ns(bs), self.x: self.ds(bs), self.xl: self.ds(4 * bs) } train_time = 0 num_epochs = 0 g_losses = [] d_losses = [] for t in range(0, max_iters): if num_epochs > hmc_epochs: use_hmc = False # eventually we want to stop using HMC as the bootstrap so we can improove upon it if t % epoch_size == 0: num_epochs += 1 self.bootstrap(sess=sess, steps=bootstrap_steps, burn_in=bootstrap_burn_in, batch_size=bootstrap_batch_size, discard_ratio=bootstrap_discard_ratio, use_hmc=use_hmc, hmc_steps=hmc_steps, hmc_step_size=hmc_step_size) if t % log_freq == 0: d_loss = sess.run(self.d_loss, feed_dict=_feed_dict(batch_size)) g_loss, v_loss = sess.run([self.g_loss, self.v_loss], feed_dict=_feed_dict(batch_size)) print( 'Iter [%d] time [%5.4f] d_loss [%.4f] g_loss [%.4f] v_loss [%.4f]' % (t, train_time, d_loss, g_loss, v_loss)) g_losses.append(g_loss) d_losses.append(d_loss) if save_path: z, sample_time = self.sample( sess, evaluate_steps + evaluate_burn_in, nice_steps, evaluate_batch_size) z = z[evaluate_burn_in:] ess = batch_means_ess(z) min_ess = np.mean(np.min(ess, axis=1), axis=0) std_ess = np.std(np.min(ess, axis=1), axis=0) acc_rate = acceptance_rate(z) with open( os.path.join(save_path['results'], 'results.txt'), 'a') as f: f.write( f"min ess at iteration {t}: {min_ess} +- {std_ess} \n" ) f.write(f"acceptance at iteration {t}: {acc_rate} \n") if plot: print('plotting') def plot2d(samples): fig, ax = plt.subplots() ax.hist2d(samples[:, 0, 0], samples[:, 0, 1], bins=400) ax.set_aspect('equal', 'box') plt.savefig( os.path.join(save_path['figs'], 'samples.png')) plt.close() plot2d(z) if (t + 1) % save_freq: saver.save(sess, os.path.join(save_path['ckpts'], 'ckpt')) start = time.time() for _ in range(0, d_iters): sess.run(self.d_train, feed_dict=_feed_dict(batch_size)) sess.run(self.g_train, feed_dict=_feed_dict(batch_size)) end = time.time() train_time += end - start g_losses = np.array(g_losses) d_losses = np.array(d_losses) losses = np.stack([g_losses, d_losses]) return losses, train_time
def train(self, sess, learning_rate=1e-3, max_iters=5000, train_samples=200, save_path='.', log_freq=100, init_temp=1.0, temp_factor=1.0, temp_rate=100, save_freq=1000, train_plot=True, evaluate_steps=500, evaluate_burn_in=100, **kwargs): global_step = tf.Variable(0., name='global_step', trainable=False) if learning_rate < 0.0: u = np.random.uniform(2.5, 7) learning_rate = np.power(u, 10) if save_path: with open(os.path.join(save_path['info'], 'params.txt'), 'a') as f: f.write("learning_rate: {learning_rate} \n") learning_rate = tf.train.exponential_decay(learning_rate, global_step, 1000, 0.96, staircase=True) optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate) train_op = optimizer.minimize(self.loss_op, global_step=global_step) sess.run(tf.global_variables_initializer()) saver = tf.train.Saver() if save_path: ckpt = tf.train.get_checkpoint_state(save_path['ckpts']) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) train_time = 0.0 losses = [] # Move attribute look ups outside loop to improve performance loss_op = self.loss_op sample_op = self.sample_op acc = self.acc x = self.x x_init = self.x_init init_dist = self.init_dist samples = init_dist(train_samples) init_samples = init_dist(train_samples) dynamics = self.dynamics for t in range(max_iters): tmp = (init_temp - 1) * (1 - t / float(max_iters)) + 1 time1 = time.time() _, loss_, samples, px_, lr_ = sess.run( [train_op, loss_op, sample_op, acc, learning_rate], { x: samples, dynamics.temperature: tmp, x_init: init_samples }) time2 = time.time() init_samples = init_dist(train_samples) train_time += time2 - time1 losses.append(loss_) if t % log_freq == 0: print( 'Time: %d Step: %d / %d, Loss: %.2e, Acceptance sample: %.2f, LR: %.5f, Temp: %.4f' % (train_time, t, max_iters, loss_, np.mean(px_), lr_, tmp)) z, _ = self.sample(sess, evaluate_steps + evaluate_burn_in, temperature=1.0) z = z[evaluate_burn_in:] ess = batch_means_ess(z) min_ess = np.mean(np.min(ess, axis=1), axis=0) std_ess = np.std(np.min(ess, axis=1), axis=0) acc_rate = acceptance_rate(z) if save_path: with open( os.path.join(save_path['results'], 'results.txt'), 'a') as f: f.write( f"min ess at iteration {t}: {min_ess} +- {std_ess} \n" ) f.write(f"acceptance at iteration {t}: {acc_rate} \n") else: print( f"min ess at iteration {t}: {min_ess} +- {std_ess} \n ", f"acceptance at iteration {t}: {acc_rate} \n") if save_path and train_plot and z.shape[2] == 2: def plot2d(samples): fig, ax = plt.subplots() ax.scatter(samples[:, 0, 0], samples[:, 0, 1]) ax.set_aspect('equal', 'box') plt.savefig( os.path.join(save_path['figs'], f"samples_{t}.png")) plt.close() plot2d(z) if save_path and (t + 1) % save_freq: saver.save(sess, os.path.join(save_path['ckpts'], 'ckpt'), global_step=global_step) return np.array(losses), train_time