def create_bound(model, xs, lengths): """Creates the bound to be evaluated.""" if config.bound == "elbo": ll_per_seq, log_weights, _ = bounds.iwae( model, xs, lengths, num_samples=1, parallel_iterations=config.parallel_iterations) elif config.bound == "iwae": ll_per_seq, log_weights, _ = bounds.iwae( model, xs, lengths, num_samples=config.num_samples, parallel_iterations=config.parallel_iterations) elif config.bound == "fivo": ll_per_seq, log_weights, resampled, _ = bounds.fivo( model, xs, lengths, num_samples=config.num_samples, resampling_criterion=smc.ess_criterion, resampling_type=config.resampling_type, random_seed=config.random_seed, parallel_iterations=config.parallel_iterations) # Compute bound scaled by number of timesteps. bound_per_t = ll_per_seq / tf.to_float(lengths) if config.bound == "fivo": return bound_per_t, log_weights, resampled else: return bound_per_t, log_weights
def create_bound(model, xs, lengths): """Creates the bound to be evaluated.""" if config.bound == "elbo": ll_per_seq, log_weights, _ = bounds.iwae( model, xs, lengths, num_samples=1, parallel_iterations=config.parallel_iterations ) elif config.bound == "iwae": ll_per_seq, log_weights, _ = bounds.iwae( model, xs, lengths, num_samples=config.num_samples, parallel_iterations=config.parallel_iterations ) elif config.bound == "fivo": ll_per_seq, log_weights, resampled, _ = bounds.fivo( model, xs, lengths, num_samples=config.num_samples, resampling_criterion=smc.ess_criterion, resampling_type=config.resampling_type, random_seed=config.random_seed, parallel_iterations=config.parallel_iterations ) # Compute bound scaled by number of timesteps. bound_per_t = ll_per_seq / tf.to_float(lengths) if config.bound == "fivo": return bound_per_t, log_weights, resampled else: return bound_per_t, log_weights
def create_loss(): """Creates the loss to be optimized. Returns: bound: A float Tensor containing the value of the bound that is being optimized. loss: A float Tensor that when differentiated yields the gradients to apply to the model. Should be optimized via gradient descent. """ inputs, targets, lengths, model, _ = create_dataset_and_model_fn( config, split="train", shuffle=True, repeat=True) # Compute lower bounds on the log likelihood if config.bound == "elbo": ll_per_seq, _, _ = bounds.iwae( model, (inputs, targets), lengths, num_samples=1, parallel_iterations=config.parallel_iterations) elif config.bound == "iwae": ll_per_seq, _, _ = bounds.iwae( model, (inputs, targets), lengths, num_samples=config.num_samples, parallel_iterations=config.parallel_iterations) elif config.bound in ("fivo", "fivo-aux"): if config.resampling_type == "relaxed": ll_per_seq, _, _, _ = bounds.fivo( model, (inputs, targets), lengths, num_samples=config.num_samples, resampling_criterion=smc.ess_criterion, resampling_type=config.resampling_type, random_seed=config.random_seed, relaxed_resampling_temperature=config. relaxed_resampling_temperature, parallel_iterations=config.parallel_iterations) else: ll_per_seq, _, _, _ = bounds.fivo( model, (inputs, targets), lengths, num_samples=config.num_samples, resampling_criterion=smc.ess_criterion, resampling_type=config.resampling_type, random_seed=config.random_seed, parallel_iterations=config.parallel_iterations) # Compute loss scaled by number of timesteps ll_per_t = tf.reduce_mean(ll_per_seq / tf.to_float(lengths)) ll_per_seq = tf.reduce_mean(ll_per_seq) tf.summary.scalar("train_ll_per_seq", ll_per_seq) tf.summary.scalar("train_ll_per_t", ll_per_t) if config.normalize_by_seq_len: return ll_per_t, -ll_per_t else: return ll_per_seq, -ll_per_seq
def create_losses(model, observations, lengths): """Creates the loss to be optimized. Args: model: A Trainable GHMM model. observations: A set of observations. lengths: The lengths of each sequence in the observations. Returns: loss: A float Tensor that when differentiated yields the gradients to apply to the model. Should be optimized via gradient descent. bound: A float Tensor containing the value of the bound that is being optimized. true_ll: The true log-likelihood of the data under the model. bound_gap: The gap between the bound and the true log-likelihood. """ # Compute lower bounds on the log likelihood. if config.bound == "elbo": ll_per_seq, _, _ = bounds.iwae( model, observations, lengths, num_samples=1, parallel_iterations=config.parallel_iterations ) elif config.bound == "iwae": ll_per_seq, _, _ = bounds.iwae( model, observations, lengths, num_samples=config.num_samples, parallel_iterations=config.parallel_iterations ) elif config.bound == "fivo": if config.resampling_type == "relaxed": ll_per_seq, _, _, _ = bounds.fivo( model, observations, lengths, num_samples=config.num_samples, resampling_criterion=smc.ess_criterion, resampling_type=config.resampling_type, relaxed_resampling_temperature=config. relaxed_resampling_temperature, random_seed=config.random_seed, parallel_iterations=config.parallel_iterations) else: ll_per_seq, _, _, _ = bounds.fivo( model, observations, lengths, num_samples=config.num_samples, resampling_criterion=smc.ess_criterion, resampling_type=config.resampling_type, random_seed=config.random_seed, parallel_iterations=config.parallel_iterations ) ll_per_t = tf.reduce_mean(ll_per_seq / tf.to_float(lengths)) # Compute the data's true likelihood under the model and the bound gap. true_ll_per_seq = model.likelihood(tf.squeeze(observations)) true_ll_per_t = tf.reduce_mean(true_ll_per_seq / tf.to_float(lengths)) bound_gap = true_ll_per_seq - ll_per_seq bound_gap = tf.reduce_mean(bound_gap/ tf.to_float(lengths)) tf.summary.scalar("train_ll_bound", ll_per_t) tf.summary.scalar("train_true_ll", true_ll_per_t) tf.summary.scalar("bound_gap", bound_gap) return -ll_per_t, ll_per_t, true_ll_per_t, bound_gap
def test_iwae(self): """A golden-value test for the IWAE bound.""" tf.set_random_seed(1234) with self.test_session() as sess: model, inputs, targets, lengths = create_vrnn(random_seed=1234) outs = bounds.iwae(model, (inputs, targets), lengths, num_samples=4, parallel_iterations=1) sess.run(tf.global_variables_initializer()) log_p_hat, weights, _ = sess.run(outs) self.assertAllClose([-23.301426, -13.64028], log_p_hat) weights_gt = np.array( [[[-3.66708851, -2.07074022, -4.91751671, -5.03293562], [-2.99690723, -3.17782736, -4.50084877, -3.48536515]], [[-6.2539978, -4.37615728, -7.43738699, -7.85044909], [-8.27518654, -6.71545124, -8.96198845, -7.05567837]], [[-9.19093227, -8.01637268, -11.64603615, -10.51128292], [-12.34527206, -11.54284477, -11.8667469, -9.69417381]], [[-12.20609856, -10.47217369, -13.66270638, -13.46115875], [-17.17656708, -16.25190353, -15.28658581, -12.33067703]], [[-16.14766312, -15.57472229, -17.47755432, -17.98189926], [-17.17656708, -16.25190353, -15.28658581, -12.33067703]], [[-20.07182884, -18.43191147, -20.1606636, -21.45263863], [-17.17656708, -16.25190353, -15.28658581, -12.33067703]], [[-24.10270691, -22.20865822, -24.14675522, -25.27248383], [-17.17656708, -16.25190353, -15.28658581, -12.33067703]]]) self.assertAllClose(weights_gt, weights)
def create_graph(): """Creates the evaluation graph. Returns: lower_bounds: A tuple of float Tensors containing the values of the 3 evidence lower bounds, summed across the batch. total_batch_length: The total number of timesteps in the batch, summed across batch examples. batch_size: The batch size. global_step: The global step the checkpoint was loaded from. """ global_step = tf.train.get_or_create_global_step() inputs, targets, lengths, model, _ = create_dataset_and_model_fn( config, split=config.split, shuffle=False, repeat=False) # Compute lower bounds on the log likelihood. elbo_ll_per_seq, _, _ = bounds.iwae( model, (inputs, targets), lengths, num_samples=1, parallel_iterations=config.parallel_iterations) iwae_ll_per_seq, _, _ = bounds.iwae( model, (inputs, targets), lengths, num_samples=config.num_samples, parallel_iterations=config.parallel_iterations) # The resampling type should only be used for training, so we ignore it. fivo_ll_per_seq, _, _, _ = bounds.fivo( model, (inputs, targets), lengths, num_samples=config.num_samples, resampling_criterion=smc.ess_criterion, random_seed=config.random_seed, parallel_iterations=config.parallel_iterations) elbo_ll = tf.reduce_sum(elbo_ll_per_seq) iwae_ll = tf.reduce_sum(iwae_ll_per_seq) fivo_ll = tf.reduce_sum(fivo_ll_per_seq) batch_size = tf.shape(lengths)[0] total_batch_length = tf.reduce_sum(lengths) return ((elbo_ll, iwae_ll, fivo_ll), total_batch_length, batch_size, global_step)
def test_elbo(self): """A golden-value test for the ELBO (the IWAE bound with num_samples=1).""" tf.set_random_seed(1234) with self.test_session() as sess: model, inputs, targets, lengths = create_vrnn(random_seed=1234) outs = bounds.iwae(model, (inputs, targets), lengths, num_samples=1, parallel_iterations=1) sess.run(tf.global_variables_initializer()) log_p_hat, _, _ = sess.run(outs) self.assertAllClose([-21.615765, -13.614225], log_p_hat)
def create_graph(): """Creates the graph to sample from the model. First, the model is conditioned on a prefix by sampling a batch of data and trimming it to prefix_length. The configured bound is used to do the conditioning. Then the final state from the conditioning is used to sample from the model. Returns: samples: A Tensor of shape [sample_length, batch_size, num_samples, data_dimension] representing samples from the model. prefixes: A Tensor of shape [prefix_length, batch_size, data_dimension] representing the prefixes the model was conditioned on. """ inputs, targets, lengths, model, mean = create_dataset_and_model_fn( config, split=config.split, shuffle=True, repeat=True) input_prefixes = inputs[:config.prefix_length] target_prefixes = targets[:config.prefix_length] prefix_lengths = tf.ones_like(lengths) * config.prefix_length if config.bound == "elbo": _, _, state = bounds.iwae(model, (input_prefixes, target_prefixes), prefix_lengths, num_samples=1) elif config.bound == "iwae": _, _, state = bounds.iwae(model, (input_prefixes, target_prefixes), prefix_lengths, num_samples=config.num_samples) elif config.bound == "fivo": _, _, _, state = bounds.fivo( model, (input_prefixes, target_prefixes), prefix_lengths, num_samples=config.num_samples, resampling_criterion=smc.ess_criterion, random_seed=config.random_seed) sample_inputs = tf.tile(inputs[config.prefix_length], [config.num_samples, 1]) samples = sample_from_model(model, state, sample_inputs, mean) return samples, target_prefixes