def SanitizedAutoCorrelationMean(x, axis, reduce_axis, max_lags=None, **kwargs): shape_arr = np.array(list(x.shape)) axes = list(sorted(set(range(len(shape_arr))) - set([reduce_axis]))) mean_shape = shape_arr[axes] if max_lags is not None: mean_shape[axis] = max_lags + 1 mean_state = fun_mc.running_mean_init(mean_shape, x.dtype) new_order = list(range(len(shape_arr))) new_order[0] = new_order[reduce_axis] new_order[reduce_axis] = 0 x = tf.transpose(x, new_order) x_arr = tf.TensorArray(x.dtype, x.shape[0]).unstack(x) mean_state, _ = fun_mc.trace( state=mean_state, fn=lambda state: fun_mc.running_mean_step( # pylint: disable=g-long-lambda state, SanitizedAutoCorrelation(x_arr.read(state.num_points), axis, max_lags=max_lags, **kwargs)), num_steps=x.shape[0], trace_fn=lambda *_: ()) return mean_state.mean
def testBasic(self): def fun(x): return x + 1., 2 * x x, _ = fun_mc.trace(state=0., fn=fun, num_steps=5) self.assertIsInstance(x, tf.Tensor)
def train_p(q, u, x_pos, step_size, opt_p): """Train P using the standard CD objective. Args: q: `ModelQ`. u: A callable representing the energy function. x_pos: A batch of positive examples. step_size: Step size to use for HMC. opt_p: A `tf.optimizer.Optimizer`. Returns: x_neg_q: Negative samples sampled from `q`. x_neg_p: Negative samples used to train `p`, possibly generated via HMC. p_accept: Acceptance rate of HMC. step_size: The new step size, possibly adapted to adjust the acceptance rate. pos_e: Mean energy of the positive samples across the batch. pos_e: Mean energy of the positive samples across the batch, after the parameter update. neg_e_q: Mean energy of `x_neg_q` across the batch. neg_e_p: Mean energy of `x_neg_p` across the batch. neg_e_p_updated: Mean energy of `x_neg_p` across the batch, after the parameter update. """ def create_momentum_sample_fn(state): sample_fn = lambda seed: tf.random.normal( # pylint: disable=g-long-lambda tf.shape(state), stddev=FLAGS.mcmc_momentum_stddev) return sample_fn _, x_neg_q, _ = q.sample_with_log_prob(FLAGS.batch_size, temp=FLAGS.q_temperature) neg_e_q = tf.reduce_mean(u(x_neg_q)) def p_log_prob(x): return -u(x) if FLAGS.use_mcmc: def log_prob_non_transformed(x): p_log_p = p_log_prob(x) return p_log_p, (x, ) # TODO(siege): Why aren't we actually using NeuTra? # def log_prob_transformed(z): # x, logdet = q.reverse(z) # p_log_p = p_log_prob(x) # return p_log_p + logdet, (x,) def kernel(hmc_state, step_size, step): """HMC kernel.""" hmc_state, hmc_extra = fun_mc.hamiltonian_monte_carlo_step( hmc_state, step_size=step_size, num_integrator_steps=FLAGS.mcmc_leapfrog_steps, momentum_sample_fn=create_momentum_sample_fn(hmc_state.state), target_log_prob_fn=log_prob_non_transformed) mean_p_accept = tf.reduce_mean( tf.exp(tf.minimum(0., hmc_extra.log_accept_ratio))) if FLAGS.mcmc_adapt_step_size: step_size = fun_mc.sign_adaptation(step_size, output=mean_p_accept, set_point=0.9) return (hmc_state, step_size, step + 1), hmc_extra hmc_state, is_accepted = fun_mc.trace( state=(fun_mc.hamiltonian_monte_carlo_init( x_neg_q, log_prob_non_transformed), step_size, 0), fn=kernel, num_steps=FLAGS.mcmc_num_steps, trace_fn=lambda _, hmc_extra: hmc_extra.is_accepted) x_neg_p = hmc_state[0].state_extra[0] step_size = hmc_state[1] p_accept = tf.reduce_mean(tf.cast(is_accepted, tf.float32)) else: x_neg_p = x_neg_q p_accept = 0.0 step_size = 0.0 with tf.GradientTape() as tape: tape.watch(u.trainable_variables) pos_e = tf.reduce_mean(u(x_pos)) neg_e_p = tf.reduce_mean(u(x_neg_p)) loss = pos_e - neg_e_p + tf.square(pos_e) * FLAGS.p_center_regularizer variables = u.trainable_variables grads = tape.gradient(loss, variables) grads_and_vars = list(zip(grads, variables)) opt_p.apply_gradients(grads_and_vars) pos_e_updated = tf.reduce_mean(u(x_pos)) neg_e_p_updated = tf.reduce_mean(u(x_neg_p)) return (x_neg_q, x_neg_p, p_accept, step_size, pos_e, pos_e_updated, neg_e_q, neg_e_p, neg_e_p_updated)