def get_kl_loss(self, means, log_vars, means_p=None, log_vars_p=None): """Get KL loss for all the predicted Gaussians.""" kl_loss = 0.0 if means_p is None: means_p = tf.unstack(tf.zeros_like(means)) if log_vars_p is None: log_vars_p = tf.unstack(tf.zeros_like(log_vars)) enumerated_inputs = enumerate(zip(means, log_vars, means_p, log_vars_p)) if self.is_training and self.hparams.stochastic_model: for i, (mean, log_var, mean_p, log_var_p) in enumerated_inputs: # Condition to compute kl divergence with learned prior kl_loss += common_layers.kl_divergence(mean, log_var, mean_p, log_var_p) tf.summary.histogram("posterior_mean_%d" % i, mean) tf.summary.histogram("posterior_log_var_%d" % i, log_var) tf.summary.histogram("prior_mean_%d" % i, mean_p) tf.summary.histogram("prior_log_var_%d" % i, log_var_p) tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss)) beta = self.get_beta(kl_loss) # information capacity from "Understanding disentangling in beta-VAE" if self.hparams.information_capacity > 0.0: kl_loss = tf.abs(kl_loss - self.hparams.information_capacity) return beta * kl_loss
def get_extra_loss(self, mean, std): """Losses in addition to the default modality losses.""" beta = self.get_beta() kl_loss = common_layers.kl_divergence(mean, std) tf.summary.histogram("posterior_mean", mean) tf.summary.histogram("posterior_std", std) tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss)) return beta * kl_loss
def get_extra_loss(self, mean, std): """Losses in addition to the default modality losses.""" kl_loss = common_layers.kl_divergence(mean, std) beta = self.get_beta(kl_loss) tf.summary.histogram("posterior_mean", mean) tf.summary.histogram("posterior_std", std) tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss)) # information capacity from "Understanding disentangling in beta-VAE" if self.hparams.information_capacity > 0.0: kl_loss = tf.abs(kl_loss - self.hparams.information_capacity) return beta * kl_loss
def get_extra_loss(self, latent_means=None, latent_stds=None, true_frames=None, gen_frames=None, beta=1.0): """Losses in addition to the default modality losses.""" kl_loss = 0.0 if self.is_training: for i, (mean, std) in enumerate(zip(latent_means, latent_stds)): kl_loss += common_layers.kl_divergence(mean, std) tf.summary.histogram("posterior_mean_%d" % i, mean) tf.summary.histogram("posterior_std_%d" % i, std) tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss)) return beta * kl_loss
def get_kl_loss(self, means, stds): """Get KL loss for all the predicted Gaussians.""" kl_loss = 0.0 if self.is_training and self.hparams.stochastic_model: for i, (mean, std) in enumerate(zip(means, stds)): kl_loss += common_layers.kl_divergence(mean, std) tf.summary.histogram("posterior_mean_%d" % i, mean) tf.summary.histogram("posterior_std_%d" % i, std) tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss)) beta = self.get_beta(kl_loss) # information capacity from "Understanding disentangling in beta-VAE" if self.hparams.information_capacity > 0.0: kl_loss = tf.abs(kl_loss - self.hparams.information_capacity) return beta * kl_loss
def get_extra_loss(self, latent_means=None, latent_stds=None, true_frames=None, gen_frames=None): """Losses in addition to the default modality losses.""" del true_frames del gen_frames kl_loss = 0.0 if self.is_training and self.hparams.stochastic_model: for i, (mean, std) in enumerate(zip(latent_means, latent_stds)): kl_loss += common_layers.kl_divergence(mean, std) tf.summary.histogram("posterior_mean_%d" % i, mean) tf.summary.histogram("posterior_std_%d" % i, std) tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss)) beta = self.get_beta(kl_loss) extra_loss = beta * kl_loss return extra_loss
def get_kl_loss(self, means, log_vars, means_p=None, log_vars_p=None): """Get KL loss for all the predicted Gaussians.""" kl_loss = 0.0 if means_p is None: means_p = tf.unstack(tf.zeros_like(means)) if log_vars_p is None: log_vars_p = tf.unstack(tf.zeros_like(log_vars)) enumerated_inputs = enumerate(zip(means, log_vars, means_p, log_vars_p)) if self.is_training and self.hparams.stochastic_model: for i, (mean, log_var, mean_p, log_var_p) in enumerated_inputs: kl_loss += common_layers.kl_divergence(mean, log_var, mean_p, log_var_p) tf.summary.histogram("posterior_mean_%d" % i, mean) tf.summary.histogram("posterior_log_var_%d" % i, log_var) tf.summary.histogram("prior_mean_%d" % i, mean_p) tf.summary.histogram("prior_log_var_%d" % i, log_var_p) tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss)) beta = self.get_beta(kl_loss) # information capacity from "Understanding disentangling in beta-VAE" if self.hparams.information_capacity > 0.0: kl_loss = tf.abs(kl_loss - self.hparams.information_capacity) return beta * kl_loss
def body(self, features): hparams = self.hparams batch_size = common_layers.shape_list(features["inputs"])[0] # Swap time and batch axes. input_frames = common_video.swap_time_and_batch_axes( features["inputs"]) target_frames = common_video.swap_time_and_batch_axes( features["targets"]) # Get actions if exist otherwise use zeros input_actions = self.get_input_if_exists( features, "input_action", batch_size, hparams.video_num_input_frames) target_actions = self.get_input_if_exists( features, "target_action", batch_size, hparams.video_num_target_frames) # Get rewards if exist otherwise use zeros input_rewards = self.get_input_if_exists( features, "input_reward", batch_size, hparams.video_num_input_frames) target_rewards = self.get_input_if_exists( features, "target_reward", batch_size, hparams.video_num_target_frames) all_actions = tf.concat([input_actions, target_actions], axis=0) all_rewards = tf.concat([input_rewards, target_rewards], axis=0) all_frames = tf.concat([input_frames, target_frames], axis=0) # Each image is being used twice, in latent tower and main tower. # This is to make sure we are using the *same* image for both, ... # ... given how TF queues work. # NOT sure if this is required at all. Doesn"t hurt though! :) all_frames = tf.identity(all_frames) gen_images, gen_rewards, latent_means, latent_stds = self.construct_model( images=all_frames, actions=all_actions, rewards=all_rewards, ) beta = self.get_beta() kl_loss = 0.0 if self.is_training: for i, (mean, std) in enumerate(zip(latent_means, latent_stds)): kl_loss += common_layers.kl_divergence(mean, std) tf.summary.histogram("posterior_mean_%d" % i, mean) tf.summary.histogram("posterior_std_%d" % i, std) tf.summary.scalar("beta", beta) tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss)) extra_loss = beta * kl_loss # Ignore the predictions from the input frames. # This is NOT the same as original paper/implementation. predictions = gen_images[hparams.video_num_input_frames - 1:] reward_pred = gen_rewards[hparams.video_num_input_frames - 1:] reward_pred = tf.squeeze(reward_pred, axis=2) # Remove undeeded dimension. # TODO(mbz): clean this up! def fix_video_dims_and_concat_on_x_axis(x): x = tf.transpose(x, [1, 3, 4, 0, 2]) x = tf.reshape(x, [batch_size, 64, 3, -1]) x = tf.transpose(x, [0, 3, 1, 2]) return x frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames) frames_pd = fix_video_dims_and_concat_on_x_axis(predictions) side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2) tf.summary.image("full_video", side_by_side_video) # Swap back time and batch axes. predictions = common_video.swap_time_and_batch_axes(predictions) reward_pred = common_video.swap_time_and_batch_axes(reward_pred) return_targets = predictions if "target_reward" in features: return_targets = { "targets": predictions, "target_reward": reward_pred } return return_targets, extra_loss
def body(self, features): hparams = self.hparams batch_size = common_layers.shape_list(features["inputs"])[0] # Swap time and batch axes. input_frames = common_video.swap_time_and_batch_axes( features["inputs"]) target_frames = common_video.swap_time_and_batch_axes( features["targets"]) # Get actions if exist otherwise use zeros input_actions = self.get_input_if_exists( features, "input_action", batch_size, hparams.video_num_input_frames) target_actions = self.get_input_if_exists( features, "target_action", batch_size, hparams.video_num_target_frames) # Get rewards if exist otherwise use zeros input_rewards = self.get_input_if_exists( features, "input_reward", batch_size, hparams.video_num_input_frames) target_rewards = self.get_input_if_exists( features, "target_reward", batch_size, hparams.video_num_target_frames) all_actions = tf.concat([input_actions, target_actions], axis=0) all_rewards = tf.concat([input_rewards, target_rewards], axis=0) all_frames = tf.concat([input_frames, target_frames], axis=0) # Each image is being used twice, in latent tower and main tower. # This is to make sure we are using the *same* image for both, ... # ... given how TF queues work. # NOT sure if this is required at all. Doesn"t hurt though! :) all_frames = tf.identity(all_frames) gen_images, gen_rewards, latent_means, latent_stds = self.construct_model( images=all_frames, actions=all_actions, rewards=all_rewards, ) step_num = tf.train.get_global_step() # TODO(mbz): what should it be if it"s undefined? if step_num is None: step_num = _LARGE_STEP_NUMBER schedule = self.hparams.latent_loss_multiplier_schedule second_stage = self.hparams.num_iterations_2nd_stage # TODO(mechcoder): Add log_annealing schedule. if schedule == "constant": beta = tf.cond(tf.greater(step_num, second_stage), lambda: self.hparams.latent_loss_multiplier, lambda: 0.0) elif schedule == "linear_anneal": # Linearly anneal beta from 0.0 to self.hparams.latent_loss_multiplier. # between self.hparams.num_iterations_2nd_stage to anneal_end. # beta = latent_loss * (1 - (global_step - 2nd_stage) / (anneal_end - 2nd_stage)) # pylint:disable=line-too-long anneal_end = self.hparams.anneal_end latent_multiplier = self.hparams.latent_loss_multiplier if anneal_end < second_stage: raise ValueError("Expected hparams.num_iterations_2nd_stage < " "hparams.anneal_end %d, got %d." % (second_stage, anneal_end)) def anneal_loss(step_num): step_num = tf.cast(step_num, dtype=tf.float32) fraction = (float(anneal_end) - step_num) / (anneal_end - second_stage) return self.hparams.latent_loss_multiplier * (1 - fraction) beta = tf.case(pred_fn_pairs={ tf.less(step_num, second_stage): lambda: 0.0, tf.greater(step_num, anneal_end): lambda: latent_multiplier }, default=lambda: anneal_loss(step_num)) kl_loss = 0.0 if self.is_training: for i, (mean, std) in enumerate(zip(latent_means, latent_stds)): kl_loss += common_layers.kl_divergence(mean, std) tf.summary.histogram("posterior_mean_%d" % i, mean) tf.summary.histogram("posterior_std_%d" % i, std) tf.summary.scalar("beta", beta) tf.summary.scalar("kl_raw", tf.reduce_mean(kl_loss)) extra_loss = beta * kl_loss # Ignore the predictions from the input frames. # This is NOT the same as original paper/implementation. predictions = gen_images[hparams.video_num_input_frames - 1:] reward_pred = gen_rewards[hparams.video_num_input_frames - 1:] reward_pred = tf.squeeze(reward_pred, axis=2) # Remove undeeded dimension. # TODO(mbz): clean this up! def fix_video_dims_and_concat_on_x_axis(x): x = tf.transpose(x, [1, 3, 4, 0, 2]) x = tf.reshape(x, [batch_size, 64, 3, -1]) x = tf.transpose(x, [0, 3, 1, 2]) return x frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames) frames_pd = fix_video_dims_and_concat_on_x_axis(predictions) side_by_side_video = tf.concat([frames_gd, frames_pd], axis=2) tf.summary.image("full_video", side_by_side_video) # Swap back time and batch axes. predictions = common_video.swap_time_and_batch_axes(predictions) reward_pred = common_video.swap_time_and_batch_axes(reward_pred) return_targets = predictions if "target_reward" in features: return_targets = { "targets": predictions, "target_reward": reward_pred } return return_targets, extra_loss