tensorflow_probability/python/mcmc/eight_schools_hmc.py
def benchmark_eight_schools_hmc(
    num_results=int(5e3),
    num_burnin_steps=int(3e3),
    num_leapfrog_steps=3,
    step_size=0.4):
  """Runs HMC on the eight-schools unnormalized posterior."""

  num_schools = 8
  treatment_effects = tf.constant(
      [28, 8, -3, 7, -1, 1, 18, 12],
      dtype=np.float32,
      name='treatment_effects')
  treatment_stddevs = tf.constant(
      [15, 10, 16, 11, 9, 11, 10, 18],
      dtype=np.float32,
      name='treatment_stddevs')

  def unnormalized_posterior_log_prob(
      avg_effect, avg_stddev, school_effects_standard):
    """Eight-schools unnormalized log posterior."""
    return eight_schools_joint_log_prob(
        treatment_effects, treatment_stddevs,
        avg_effect, avg_stddev, school_effects_standard)

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

  def computation():
    """The benchmark computation."""
    _, kernel_results = sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=(
            tf.zeros([], name='init_avg_effect'),
            tf.zeros([], name='init_avg_stddev'),
            tf.ones([num_schools], name='init_school_effects_standard'),
        ),
        kernel=tfp.mcmc.HamiltonianMonteCarlo(
            target_log_prob_fn=unnormalized_posterior_log_prob,
            step_size=step_size,
            num_leapfrog_steps=num_leapfrog_steps))

    return kernel_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(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_eight_schools_hmc(
    num_results=int(5e3),
    num_burnin_steps=int(3e3),
    num_leapfrog_steps=3,
    step_size=0.4):
  num_schools = 8
  treatment_effects = tf.constant(
      [28, 8, -3, 7, -1, 1, 18, 12],
      dtype=np.float32,
      name='treatment_effects')
  treatment_stddevs = tf.constant(
      [15, 10, 16, 11, 9, 11, 10, 18],
      dtype=np.float32,
      name='treatment_stddevs')

  def unnormalized_posterior_log_prob(
      avg_effect, avg_stddev, school_effects_standard):
    """Eight-schools unnormalized log posterior."""
    return eight_schools_joint_log_prob(
        treatment_effects, treatment_stddevs,
        avg_effect, avg_stddev, school_effects_standard)

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

  def computation():
    """The benchmark computation."""
    _, kernel_results = sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=(
            tf.zeros([], name='init_avg_effect'),
            tf.zeros([], name='init_avg_stddev'),
            tf.ones([num_schools], name='init_school_effects_standard'),
        ),
        kernel=tfp.mcmc.HamiltonianMonteCarlo(
            target_log_prob_fn=unnormalized_posterior_log_prob,
            step_size=step_size,
            num_leapfrog_steps=num_leapfrog_steps))

    return kernel_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(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_eight_schools_hmc ( num_results = int ( 5e3 ) , num_burnin_steps = int ( 3e3 ) , num_leapfrog_steps = 3 , step_size = 0.4 ) : num_schools = 8 treatment_effects = tf . constant ( [ 28 , 8 , - 3 , 7 , - 1 , 1 , 18 , 12 ] , dtype = np . float32 , name = 'treatment_effects' ) treatment_stddevs = tf . constant ( [ 15 , 10 , 16 , 11 , 9 , 11 , 10 , 18 ] , dtype = np . float32 , name = 'treatment_stddevs' ) def unnormalized_posterior_log_prob ( avg_effect , avg_stddev , school_effects_standard ) : """Eight-schools unnormalized log posterior.""" return eight_schools_joint_log_prob ( treatment_effects , treatment_stddevs , avg_effect , avg_stddev , school_effects_standard ) if tf . executing_eagerly ( ) : sample_chain = tf . function ( tfp . mcmc . sample_chain ) else : sample_chain = tfp . mcmc . sample_chain def computation ( ) : """The benchmark computation.""" _ , kernel_results = sample_chain ( num_results = num_results , num_burnin_steps = num_burnin_steps , current_state = ( tf . zeros ( [ ] , name = 'init_avg_effect' ) , tf . zeros ( [ ] , name = 'init_avg_stddev' ) , tf . ones ( [ num_schools ] , name = 'init_school_effects_standard' ) , ) , kernel = tfp . mcmc . HamiltonianMonteCarlo ( target_log_prob_fn = unnormalized_posterior_log_prob , step_size = step_size , num_leapfrog_steps = num_leapfrog_steps ) ) return kernel_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 ( 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 eight-schools unnormalized posterior.
Runs HMC on the eight - schools unnormalized posterior .