(_)
| 34 | |
| 35 | |
| 36 | def main(_): |
| 37 | tf.enable_eager_execution() |
| 38 | global_step = tf.train.get_or_create_global_step() |
| 39 | global_step.assign(1) |
| 40 | |
| 41 | energy_fn, mean, covar = { |
| 42 | "scg": l2hmc.get_scg_energy_fn(), |
| 43 | "rw": l2hmc.get_rw_energy_fn() |
| 44 | }[FLAGS.energy_fn] |
| 45 | |
| 46 | x_dim = 2 |
| 47 | train_iters = 5000 |
| 48 | eval_iters = 2000 |
| 49 | eps = 0.1 |
| 50 | n_steps = 10 # Chain length |
| 51 | n_samples = 200 |
| 52 | record_loss_every = 100 |
| 53 | |
| 54 | dynamics = l2hmc.Dynamics( |
| 55 | x_dim=x_dim, minus_loglikelihood_fn=energy_fn, n_steps=n_steps, eps=eps) |
| 56 | learning_rate = tf.train.exponential_decay( |
| 57 | 1e-3, global_step, 1000, 0.96, staircase=True) |
| 58 | optimizer = tf.train.AdamOptimizer(learning_rate) |
| 59 | checkpointer = tf.train.Checkpoint( |
| 60 | optimizer=optimizer, dynamics=dynamics, global_step=global_step) |
| 61 | |
| 62 | if FLAGS.train_dir: |
| 63 | summary_writer = tf.contrib.summary.create_file_writer(FLAGS.train_dir) |
| 64 | if FLAGS.restore: |
| 65 | latest_path = tf.train.latest_checkpoint(FLAGS.train_dir) |
| 66 | checkpointer.restore(latest_path) |
| 67 | print("Restored latest checkpoint at path:\"{}\" ".format(latest_path)) |
| 68 | sys.stdout.flush() |
| 69 | |
| 70 | if not FLAGS.restore: |
| 71 | # Training |
| 72 | if FLAGS.use_defun: |
| 73 | # Use `tfe.deun` to boost performance when there are lots of small ops |
| 74 | loss_fn = tfe.function(l2hmc.compute_loss) |
| 75 | else: |
| 76 | loss_fn = l2hmc.compute_loss |
| 77 | |
| 78 | samples = tf.random_normal(shape=[n_samples, x_dim]) |
| 79 | for i in range(1, train_iters + 1): |
| 80 | loss, samples, accept_prob = train_one_iter( |
| 81 | dynamics, |
| 82 | samples, |
| 83 | optimizer, |
| 84 | loss_fn=loss_fn, |
| 85 | global_step=global_step) |
| 86 | |
| 87 | if i % record_loss_every == 0: |
| 88 | print("Iteration {}, loss {:.4f}, x_accept_prob {:.4f}".format( |
| 89 | i, loss.numpy(), |
| 90 | accept_prob.numpy().mean())) |
| 91 | if FLAGS.train_dir: |
| 92 | with summary_writer.as_default(): |
| 93 | with tf.contrib.summary.always_record_summaries(): |
nothing calls this directly
no test coverage detected