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