def benchmarkEagerL2hmc(self):
"""Benchmark Eager performance."""
hparams = get_default_hparams()
dynamics = l2hmc.Dynamics(x_dim=hparams.x_dim,
loglikelihood_fn=l2hmc.get_scg_energy_fn(),
n_steps=hparams.n_steps,
eps=hparams.eps)
# TODO(lxuechen): Add learning rate decay
optimizer = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate)
# Warmup to reduce initialization effect when timing
l2hmc.warmup(dynamics, optimizer, n_iters=hparams.n_warmup_iters)
# Time
start_time = time.time()
l2hmc.fit(dynamics,
optimizer,
n_samples=hparams.n_samples,
n_iters=hparams.n_iters)
wall_time = time.time() - start_time
examples_per_sec = hparams.n_samples / wall_time
self.report_benchmark(name="eager_train_%s" %
("gpu" if tfe.num_gpus() > 0 else "cpu"),
iters=hparams.n_iters,
extras={"examples_per_sec": examples_per_sec},
wall_time=wall_time)
def benchmark_graph(self):
"""Benchmark Graph performance."""
hparams = get_default_hparams()
tf.reset_default_graph()
with tf.Graph().as_default():
energy_fn, _, _ = l2hmc.get_scg_energy_fn()
dynamics = l2hmc.Dynamics(x_dim=hparams.x_dim,
minus_loglikelihood_fn=energy_fn,
n_steps=hparams.n_steps,
eps=hparams.eps)
x = tf.placeholder(tf.float32, shape=[None, hparams.x_dim])
loss, x_out, _ = l2hmc.compute_loss(dynamics, x)
global_step = tf.Variable(0., name="global_step", trainable=False)
learning_rate = tf.train.exponential_decay(hparams.learning_rate,
global_step,
1000,
0.96,
staircase=True)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
train_op = optimizer.minimize(loss, global_step=global_step)
# Single thread; fairer comparison against eager
session_conf = tf.ConfigProto(intra_op_parallelism_threads=1,
inter_op_parallelism_threads=1)
with tf.Session(config=session_conf) as sess:
sess.run(tf.global_variables_initializer())
# Warmup to reduce initialization effect when timing
samples = npr.normal(size=[hparams.n_samples, hparams.x_dim])
for _ in range(hparams.n_warmup_iters):
_, _, _, _ = sess.run(
[x_out, loss, train_op, learning_rate],
feed_dict={x: samples})
# Training
start_time = time.time()
for i in range(hparams.n_iters):
samples, loss_np, _, _ = sess.run(
[x_out, loss, train_op, learning_rate],
feed_dict={x: samples})
print("Iteration %d: loss %.4f" % (i, loss_np))
wall_time = time.time() - start_time
examples_per_sec = hparams.n_samples / wall_time
self.report_benchmark(
name="graph_train_%s" %
("gpu" if tf.test.is_gpu_available() else "cpu"),
iters=hparams.n_iters,
extras={"examples_per_sec": examples_per_sec},
wall_time=wall_time)
def test_apply_transition(self):
"""Testing function `Dynamics.apply_transition` in graph and eager mode."""
# Eager mode testing
hparams = get_default_hparams()
energy_fn, _, _ = l2hmc.get_scg_energy_fn()
dynamics = l2hmc.Dynamics(x_dim=hparams.x_dim,
minus_loglikelihood_fn=energy_fn,
n_steps=hparams.n_steps,
eps=hparams.eps)
samples = tf.random_normal(shape=[hparams.n_samples, hparams.x_dim])
x_, v_, x_accept_prob, x_out = dynamics.apply_transition(samples)
self.assertEqual(x_.shape, v_.shape)
self.assertEqual(x_out.shape, samples.shape)
self.assertEqual(x_.shape, x_out.shape)
self.assertEqual(x_accept_prob.shape, (hparams.n_samples, ))
# Graph mode testing
with tf.Graph().as_default():
energy_fn, _, _ = l2hmc.get_scg_energy_fn()
dynamics = l2hmc.Dynamics(x_dim=hparams.x_dim,
minus_loglikelihood_fn=energy_fn,
n_steps=hparams.n_steps,
eps=hparams.eps)
x = tf.placeholder(tf.float32, shape=[None, hparams.x_dim])
x_, v_, x_accept_prob, x_out = dynamics.apply_transition(x)
samples = npr.normal(size=[hparams.n_samples, hparams.x_dim])
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
np_x_, np_v_, np_x_accept_prob, np_x_out = sess.run(
[x_, v_, x_accept_prob, x_out], feed_dict={x: samples})
self.assertEqual(np_x_.shape, np_v_.shape)
self.assertEqual(samples.shape, np_x_out.shape)
self.assertEqual(np_x_.shape, np_x_out.shape)
self.assertEqual(np_x_accept_prob.shape, (hparams.n_samples, ))
def testComputeLoss(self):
"""Testing function l2hmc.compute_loss in both graph and eager mode."""
# Eager mode testing
hparams = get_default_hparams()
dynamics = l2hmc.Dynamics(x_dim=hparams.x_dim,
loglikelihood_fn=l2hmc.get_scg_energy_fn(),
n_steps=hparams.n_steps,
eps=hparams.eps)
samples = tf.random_normal(shape=[hparams.n_samples, hparams.x_dim])
loss, x_out = l2hmc.compute_loss(samples, dynamics)
# Check shape and numerical stability
self.assertEqual(x_out.shape, samples.shape)
self.assertEqual(loss.shape, [])
self.assertAllClose(loss.numpy(), loss.numpy(), rtol=1e-5)
# Graph mode testing
with tf.Graph().as_default():
dynamics = l2hmc.Dynamics(
x_dim=hparams.x_dim,
loglikelihood_fn=l2hmc.get_scg_energy_fn(),
n_steps=hparams.n_steps,
eps=hparams.eps)
x = tf.placeholder(tf.float32, shape=[None, hparams.x_dim])
loss, x_out = l2hmc.compute_loss(x, dynamics)
samples = npr.normal(size=[hparams.n_samples, hparams.x_dim])
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
loss_np, x_out_np = sess.run([loss, x_out],
feed_dict={x: samples})
# Check shape and numerical stability
self.assertEqual(x_out_np.shape, samples.shape)
self.assertEqual(loss_np.shape, ())
self.assertAllClose(loss_np, loss_np, rtol=1e-5)
def benchmarkGraphL2hmc(self):
"""Benchmark Graph performance."""
hparams = get_default_hparams()
with tf.Graph().as_default():
dynamics = l2hmc.Dynamics(
x_dim=hparams.x_dim,
loglikelihood_fn=l2hmc.get_scg_energy_fn(),
n_steps=hparams.n_steps,
eps=hparams.eps)
x = tf.placeholder(tf.float32, shape=[None, hparams.x_dim])
loss, x_out = l2hmc.compute_loss(x, dynamics)
global_step = tf.Variable(0., name="global_step", trainable=False)
learning_rate = tf.train.exponential_decay(hparams.learning_rate,
global_step,
1000,
0.96,
staircase=True)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
train_op = optimizer.minimize(loss, global_step=global_step)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
# Warmup to reduce initialization effect when timing
samples = npr.normal(size=[hparams.n_samples, hparams.x_dim])
for _ in range(hparams.n_warmup_iters):
samples, _, _, _ = sess.run(
[x_out, loss, train_op, learning_rate],
feed_dict={x: samples})
# Time
start_time = time.time()
for _ in range(hparams.n_iters):
samples, _, _, _ = sess.run(
[x_out, loss, train_op, learning_rate],
feed_dict={x: samples})
wall_time = time.time() - start_time
examples_per_sec = hparams.n_samples / wall_time
self.report_benchmark(
name="graph_train_%s" %
("gpu" if tf.test.is_gpu_available() else "cpu"),
iters=hparams.n_iters,
extras={"examples_per_sec": examples_per_sec},
wall_time=wall_time)
def benchmark_graph(self):
"""Benchmark Graph performance."""
hparams = get_default_hparams()
tf.enable_resource_variables()
for sample_size in [10, 25, 50, 100, 200]:
hparams.n_samples = sample_size
tf.reset_default_graph()
with tf.Graph().as_default():
energy_fn, _, _ = l2hmc.get_scg_energy_fn()
x = tf.random_normal([hparams.n_samples, hparams.x_dim],
dtype=tf.float32)
dynamics = l2hmc.Dynamics(x_dim=hparams.x_dim,
minus_loglikelihood_fn=energy_fn,
n_steps=hparams.n_steps,
eps=hparams.eps)
loss, _, _ = l2hmc.compute_loss(dynamics, x)
optimizer = tf.train.AdamOptimizer(
learning_rate=hparams.learning_rate)
train_op, loss, _ = graph_step(dynamics, optimizer, x)
# Single thread; fairer comparison against eager
session_conf = tf.ConfigProto(inter_op_parallelism_threads=1)
with tf.Session(config=session_conf) as sess:
sess.run(tf.global_variables_initializer())
# Warmup to reduce initialization effect when timing
for _ in range(hparams.n_warmup_iters):
_, _ = sess.run([train_op, loss])
# Training
start_time = time.time()
for i in range(hparams.n_iters):
_, loss_np = sess.run([train_op, loss])
print("Iteration %d: loss %.4f" % (i, loss_np))
wall_time = (time.time() - start_time) / hparams.n_iters
examples_per_sec = hparams.n_samples / wall_time
self.report_benchmark(
name="graph_train_%s_%d" %
("gpu" if tf.test.is_gpu_available() else "cpu",
sample_size),
iters=hparams.n_iters,
extras={"examples_per_sec": examples_per_sec},
wall_time=wall_time)
def _benchmark_eager(self, defun=False):
"""Benchmark Eager performance."""
hparams = get_default_hparams()
for sample_size in [10, 25, 50, 100, 200]:
hparams.n_samples = sample_size
energy_fn, _, _ = l2hmc.get_scg_energy_fn()
dynamics = l2hmc.Dynamics(x_dim=hparams.x_dim,
minus_loglikelihood_fn=energy_fn,
n_steps=hparams.n_steps,
eps=hparams.eps)
optimizer = tf.train.AdamOptimizer(
learning_rate=hparams.learning_rate)
step_fn = tfe.defun(step) if defun else step
# Warmup to reduce initialization effect when timing
warmup(dynamics,
optimizer,
n_iters=hparams.n_warmup_iters,
n_samples=hparams.n_samples,
step_fn=step_fn)
# Training
samples = tf.random_normal(
shape=[hparams.n_samples, hparams.x_dim], dtype=tf.float32)
start_time = time.time()
fit(dynamics,
samples,
optimizer,
step_fn=step_fn,
n_iters=hparams.n_iters)
wall_time = (time.time() - start_time) / hparams.n_iters
examples_per_sec = hparams.n_samples / wall_time
self.report_benchmark(
name="eager_train_%s%s_%d" %
("gpu" if tf.test.is_gpu_available() else "cpu",
"_defun" if defun else "", sample_size),
iters=hparams.n_iters,
extras={"examples_per_sec": examples_per_sec},
wall_time=wall_time)
del dynamics
def _benchmark_eager(self, defun=False):
"""Benchmark Eager performance."""
hparams = get_default_hparams()
energy_fn = self._get_energy_fn()
dynamics = l2hmc.Dynamics(x_dim=hparams.x_dim,
loglikelihood_fn=energy_fn,
n_steps=hparams.n_steps,
eps=hparams.eps)
optimizer = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate)
loss_fn = tfe.defun(compute_loss) if defun else compute_loss
# Warmup to reduce initialization effect when timing
warmup(dynamics,
optimizer,
n_iters=hparams.n_warmup_iters,
loss_fn=loss_fn)
# Training
samples = tf.random_normal(shape=[hparams.n_samples, hparams.x_dim],
dtype=tf.float32)
start_time = time.time()
fit(dynamics,
samples,
optimizer,
loss_fn=loss_fn,
n_iters=hparams.n_iters,
decay_lr=True)
wall_time = time.time() - start_time
examples_per_sec = hparams.n_samples / wall_time
self.report_benchmark(name="eager_train_%s%s" %
("gpu" if tf.test.is_gpu_available() else "cpu",
"_defun" if defun else ""),
iters=hparams.n_iters,
extras={"examples_per_sec": examples_per_sec},
wall_time=wall_time)
del dynamics
del loss_fn
def main(_):
tf.enable_eager_execution()
global_step = tf.train.get_or_create_global_step()
global_step.assign(1)
energy_fn, mean, covar = {
"scg": l2hmc.get_scg_energy_fn(),
"rw": l2hmc.get_rw_energy_fn()
}[FLAGS.energy_fn]
x_dim = 2
train_iters = 5000
eval_iters = 2000
eps = 0.1
n_steps = 10 # Chain length
n_samples = 200
record_loss_every = 100
dynamics = l2hmc.Dynamics(x_dim=x_dim,
minus_loglikelihood_fn=energy_fn,
n_steps=n_steps,
eps=eps)
learning_rate = tf.train.exponential_decay(1e-3,
global_step,
1000,
0.96,
staircase=True)
optimizer = tf.train.AdamOptimizer(learning_rate)
checkpointer = tf.train.Checkpoint(optimizer=optimizer,
dynamics=dynamics,
global_step=global_step)
if FLAGS.train_dir:
summary_writer = tf.contrib.summary.create_file_writer(FLAGS.train_dir)
if FLAGS.restore:
latest_path = tf.train.latest_checkpoint(FLAGS.train_dir)
checkpointer.restore(latest_path)
print("Restored latest checkpoint at path:\"{}\" ".format(
latest_path))
sys.stdout.flush()
if not FLAGS.restore:
# Training
if FLAGS.use_defun:
# Use `tfe.deun` to boost performance when there are lots of small ops
loss_fn = tfe.defun(l2hmc.compute_loss)
else:
loss_fn = l2hmc.compute_loss
samples = tf.random_normal(shape=[n_samples, x_dim])
for i in range(1, train_iters + 1):
loss, samples, accept_prob = train_one_iter(
dynamics,
samples,
optimizer,
loss_fn=loss_fn,
global_step=global_step)
if i % record_loss_every == 0:
print("Iteration {}, loss {:.4f}, x_accept_prob {:.4f}".format(
i, loss.numpy(),
accept_prob.numpy().mean()))
if FLAGS.train_dir:
with summary_writer.as_default():
with tf.contrib.summary.always_record_summaries():
tf.contrib.summary.scalar("Training loss",
loss,
step=global_step)
print("Training complete.")
sys.stdout.flush()
if FLAGS.train_dir:
saved_path = checkpointer.save(
file_prefix=os.path.join(FLAGS.train_dir, "ckpt"))
print("Saved checkpoint at path: \"{}\" ".format(saved_path))
sys.stdout.flush()
# Evaluation
if FLAGS.use_defun:
# Use tfe.deun to boost performance when there are lots of small ops
apply_transition = tfe.defun(dynamics.apply_transition)
else:
apply_transition = dynamics.apply_transition
samples = tf.random_normal(shape=[n_samples, x_dim])
samples_history = []
for i in range(eval_iters):
samples_history.append(samples.numpy())
_, _, _, samples = apply_transition(samples)
samples_history = np.array(samples_history)
print("Sampling complete.")
sys.stdout.flush()
# Mean and covariance of target distribution
mean = mean.numpy()
covar = covar.numpy()
ac_spectrum = compute_ac_spectrum(samples_history, mean, covar)
print("First 25 entries of the auto-correlation spectrum: {}".format(
ac_spectrum[:25]))
ess = compute_ess(ac_spectrum)
print("Effective sample size per Metropolis-Hastings step: {}".format(ess))
sys.stdout.flush()
if FLAGS.train_dir:
# Plot autocorrelation spectrum in tensorboard
plot_step = tfe.Variable(1, trainable=False, dtype=tf.int64)
for ac in ac_spectrum:
with summary_writer.as_default():
with tf.contrib.summary.always_record_summaries():
tf.contrib.summary.scalar("Autocorrelation",
ac,
step=plot_step)
plot_step.assign(plot_step + n_steps)
if HAS_MATPLOTLIB:
# Choose a single chain and plot the trajectory
single_chain = samples_history[:, 0, :]
xs = single_chain[:100, 0]
ys = single_chain[:100, 1]
plt.figure()
plt.plot(xs, ys, color="orange", marker="o",
alpha=0.6) # Trained chain
plt.savefig(os.path.join(FLAGS.train_dir, "single_chain.png"))
