tensorflow_probability/python/mcmc/text_messages_hmc.py
def benchmark_text_messages_hmc(
    num_results=int(3e3),
    num_burnin_steps=int(3e3),
    num_leapfrog_steps=3):
  """Runs HMC on the text-messages unnormalized posterior."""

  if not tf.executing_eagerly():
    tf.compat.v1.reset_default_graph()

  # Build a static, pretend dataset.
  count_data = tf.cast(
      tf.concat(
          [tfd.Poisson(rate=15.).sample(43),
           tfd.Poisson(rate=25.).sample(31)],
          axis=0),
      dtype=tf.float32)
  if tf.executing_eagerly():
    count_data = count_data.numpy()
  else:
    with tf.compat.v1.Session():
      count_data = count_data.eval()

  # Define a closure over our joint_log_prob.
  def unnormalized_log_posterior(lambda1, lambda2, tau):
    return text_messages_joint_log_prob(count_data, lambda1, lambda2, tau)

  if tf.executing_eagerly():
    sample_chain = tf.function(tfp.mcmc.sample_chain)
  else:
    sample_chain = tfp.mcmc.sample_chain

  # Initialize the step_size. (It will be automatically adapted.)
  step_size = tf.compat.v2.Variable(
      name='step_size',
      initial_value=tf.constant(0.05, dtype=tf.float32),
      trainable=False)

  def computation():
    """The benchmark computation."""

    initial_chain_state = [
        tf.constant(count_data.mean(), name='init_lambda1'),
        tf.constant(count_data.mean(), name='init_lambda2'),
        tf.constant(0.5, name='init_tau'),
    ]

    unconstraining_bijectors = [
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Sigmoid(),   # Maps [0,1] to R.
    ]

    _, kernel_results = sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=initial_chain_state,
        kernel=tfp.mcmc.TransformedTransitionKernel(
            inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
                target_log_prob_fn=unnormalized_log_posterior,
                num_leapfrog_steps=num_leapfrog_steps,
                step_size=step_size,
                step_size_update_fn=
                tfp.mcmc.make_simple_step_size_update_policy(num_burnin_steps),
                state_gradients_are_stopped=True),
            bijector=unconstraining_bijectors))

    return kernel_results.inner_results.is_accepted

  # Let's force evaluation of graph to ensure build time is not part of our time
  # trial.
  is_accepted_tensor = computation()
  if not tf.executing_eagerly():
    session = tf.compat.v1.Session()
    session.run(tf.compat.v1.global_variables_initializer())
    session.run(is_accepted_tensor)

  start_time = time.time()
  if tf.executing_eagerly():
    is_accepted = computation()
  else:
    is_accepted = session.run(is_accepted_tensor)
  wall_time = time.time() - start_time

  num_accepted = np.sum(is_accepted)
  acceptance_rate = np.float32(num_accepted) / np.float32(num_results)

  return dict(
      iters=(num_results + num_burnin_steps) * num_leapfrog_steps,
      extras={'acceptance_rate': acceptance_rate},
      wall_time=wall_time)
def benchmark_text_messages_hmc(
    num_results=int(3e3),
    num_burnin_steps=int(3e3),
    num_leapfrog_steps=3):
  if not tf.executing_eagerly():
    tf.compat.v1.reset_default_graph()

  # Build a static, pretend dataset.
  count_data = tf.cast(
      tf.concat(
          [tfd.Poisson(rate=15.).sample(43),
           tfd.Poisson(rate=25.).sample(31)],
          axis=0),
      dtype=tf.float32)
  if tf.executing_eagerly():
    count_data = count_data.numpy()
  else:
    with tf.compat.v1.Session():
      count_data = count_data.eval()

  # Define a closure over our joint_log_prob.
  def unnormalized_log_posterior(lambda1, lambda2, tau):
    return text_messages_joint_log_prob(count_data, lambda1, lambda2, tau)

  if tf.executing_eagerly():
    sample_chain = tf.function(tfp.mcmc.sample_chain)
  else:
    sample_chain = tfp.mcmc.sample_chain

  # Initialize the step_size. (It will be automatically adapted.)
  step_size = tf.compat.v2.Variable(
      name='step_size',
      initial_value=tf.constant(0.05, dtype=tf.float32),
      trainable=False)

  def computation():
    """The benchmark computation."""

    initial_chain_state = [
        tf.constant(count_data.mean(), name='init_lambda1'),
        tf.constant(count_data.mean(), name='init_lambda2'),
        tf.constant(0.5, name='init_tau'),
    ]

    unconstraining_bijectors = [
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Sigmoid(),   # Maps [0,1] to R.
    ]

    _, kernel_results = sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=initial_chain_state,
        kernel=tfp.mcmc.TransformedTransitionKernel(
            inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
                target_log_prob_fn=unnormalized_log_posterior,
                num_leapfrog_steps=num_leapfrog_steps,
                step_size=step_size,
                step_size_update_fn=
                tfp.mcmc.make_simple_step_size_update_policy(num_burnin_steps),
                state_gradients_are_stopped=True),
            bijector=unconstraining_bijectors))

    return kernel_results.inner_results.is_accepted

  # Let's force evaluation of graph to ensure build time is not part of our time
  # trial.
  is_accepted_tensor = computation()
  if not tf.executing_eagerly():
    session = tf.compat.v1.Session()
    session.run(tf.compat.v1.global_variables_initializer())
    session.run(is_accepted_tensor)

  start_time = time.time()
  if tf.executing_eagerly():
    is_accepted = computation()
  else:
    is_accepted = session.run(is_accepted_tensor)
  wall_time = time.time() - start_time

  num_accepted = np.sum(is_accepted)
  acceptance_rate = np.float32(num_accepted) / np.float32(num_results)

  return dict(
      iters=(num_results + num_burnin_steps) * num_leapfrog_steps,
      extras={'acceptance_rate': acceptance_rate},
      wall_time=wall_time)
def benchmark_text_messages_hmc ( num_results = int ( 3e3 ) , num_burnin_steps = int ( 3e3 ) , num_leapfrog_steps = 3 ) : if not tf . executing_eagerly ( ) : tf . compat . v1 . reset_default_graph ( ) # Build a static, pretend dataset. count_data = tf . cast ( tf . concat ( [ tfd . Poisson ( rate = 15. ) . sample ( 43 ) , tfd . Poisson ( rate = 25. ) . sample ( 31 ) ] , axis = 0 ) , dtype = tf . float32 ) if tf . executing_eagerly ( ) : count_data = count_data . numpy ( ) else : with tf . compat . v1 . Session ( ) : count_data = count_data . eval ( ) # Define a closure over our joint_log_prob. def unnormalized_log_posterior ( lambda1 , lambda2 , tau ) : return text_messages_joint_log_prob ( count_data , lambda1 , lambda2 , tau ) if tf . executing_eagerly ( ) : sample_chain = tf . function ( tfp . mcmc . sample_chain ) else : sample_chain = tfp . mcmc . sample_chain # Initialize the step_size. (It will be automatically adapted.) step_size = tf . compat . v2 . Variable ( name = 'step_size' , initial_value = tf . constant ( 0.05 , dtype = tf . float32 ) , trainable = False ) def computation ( ) : """The benchmark computation.""" initial_chain_state = [ tf . constant ( count_data . mean ( ) , name = 'init_lambda1' ) , tf . constant ( count_data . mean ( ) , name = 'init_lambda2' ) , tf . constant ( 0.5 , name = 'init_tau' ) , ] unconstraining_bijectors = [ tfp . bijectors . Exp ( ) , # Maps a positive real to R. tfp . bijectors . Exp ( ) , # Maps a positive real to R. tfp . bijectors . Sigmoid ( ) , # Maps [0,1] to R. ] _ , kernel_results = sample_chain ( num_results = num_results , num_burnin_steps = num_burnin_steps , current_state = initial_chain_state , kernel = tfp . mcmc . TransformedTransitionKernel ( inner_kernel = tfp . mcmc . HamiltonianMonteCarlo ( target_log_prob_fn = unnormalized_log_posterior , num_leapfrog_steps = num_leapfrog_steps , step_size = step_size , step_size_update_fn = tfp . mcmc . make_simple_step_size_update_policy ( num_burnin_steps ) , state_gradients_are_stopped = True ) , bijector = unconstraining_bijectors ) ) return kernel_results . inner_results . is_accepted # Let's force evaluation of graph to ensure build time is not part of our time # trial. is_accepted_tensor = computation ( ) if not tf . executing_eagerly ( ) : session = tf . compat . v1 . Session ( ) session . run ( tf . compat . v1 . global_variables_initializer ( ) ) session . run ( is_accepted_tensor ) start_time = time . time ( ) if tf . executing_eagerly ( ) : is_accepted = computation ( ) else : is_accepted = session . run ( is_accepted_tensor ) wall_time = time . time ( ) - start_time num_accepted = np . sum ( is_accepted ) acceptance_rate = np . float32 ( num_accepted ) / np . float32 ( num_results ) return dict ( iters = ( num_results + num_burnin_steps ) * num_leapfrog_steps , extras = { 'acceptance_rate' : acceptance_rate } , wall_time = wall_time )
Runs HMC on the text-messages unnormalized posterior.
Runs HMC on the text - messages unnormalized posterior .