def model_fn(self, features, labels, mode, params): """TPUEstimator compatible model function.""" del labels is_training = (mode == tf.estimator.ModeKeys.TRAIN) data_shape = features.get_shape().as_list()[1:] z_mean, z_logvar = self.gaussian_encoder(features, is_training=is_training) z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar) reconstructions = self.decode(z_sampled, data_shape, is_training) per_sample_loss = losses.make_reconstruction_loss( features, reconstructions) reconstruction_loss = tf.reduce_mean(per_sample_loss) kl_loss = compute_gaussian_kl(z_mean, z_logvar) regularizer = self.regularizer(kl_loss, z_mean, z_logvar, z_sampled) loss = tf.add(reconstruction_loss, regularizer, name="loss") elbo = tf.add(reconstruction_loss, kl_loss, name="elbo") additional_logging = self.get_additional_logging() if mode == tf.estimator.ModeKeys.TRAIN: optimizer = optimizers.make_vae_optimizer() self._optimizer = optimizer update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) train_op = optimizer.minimize( loss=loss, global_step=tf.train.get_global_step()) train_op = tf.group( [train_op, update_ops, *self.get_additional_ops()]) tf.summary.scalar("reconstruction_loss", reconstruction_loss) tf.summary.scalar("elbo", -elbo) for log_name, log_val in additional_logging.items(): tf.summary.scalar(log_name, log_val) logging_hook = tf.train.LoggingTensorHook( { "loss": loss, "reconstruction_loss": reconstruction_loss, "elbo": -elbo, **additional_logging, }, every_n_iter=100) return tf.contrib.tpu.TPUEstimatorSpec( mode=mode, loss=loss, train_op=train_op, training_hooks=[logging_hook]) elif mode == tf.estimator.ModeKeys.EVAL: return tf.contrib.tpu.TPUEstimatorSpec( mode=mode, loss=loss, eval_metrics=(make_metric_fn("reconstruction_loss", "elbo", "regularizer", "kl_loss", *additional_logging.keys()), [ reconstruction_loss, -elbo, regularizer, kl_loss, *additional_logging.values() ])) else: raise NotImplementedError("Eval mode not supported.")
def model_fn(self, features, labels, mode, params): """TPUEstimator compatible model function.""" del labels is_training = (mode == tf.estimator.ModeKeys.TRAIN) data_shape = features.get_shape().as_list()[1:] batch_size = tf.shape(features)[0] z_mean, z_logvar = self.gaussian_encoder(features, is_training=is_training) z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar) # z_sampled_sum = z_sampled[:batch_size // 2] + \ # z_sampled[batch_size // 2:] # z_sampled_all = tf.concat([z_sampled, z_sampled_sum], axis=0) z_sampled_all = z_sampled reconstructions, group_feats_G, lie_alg_basis = self.decode_with_gfeats( z_sampled_all, data_shape, is_training) per_sample_loss = losses.make_reconstruction_loss( features, reconstructions[:batch_size]) reconstruction_loss = tf.reduce_mean(per_sample_loss) kl_loss = compute_gaussian_kl(z_mean, z_logvar) regularizer = self.regularizer(kl_loss, z_mean, z_logvar, z_sampled, group_feats_G, lie_alg_basis, batch_size) loss = tf.add(reconstruction_loss, regularizer, name="loss") elbo = tf.add(reconstruction_loss, kl_loss, name="elbo") if mode == tf.estimator.ModeKeys.TRAIN: optimizer = optimizers.make_vae_optimizer() update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) train_op = optimizer.minimize( loss=loss, global_step=tf.train.get_global_step()) train_op = tf.group([train_op, update_ops]) tf.summary.scalar("reconstruction_loss", reconstruction_loss) tf.summary.scalar("elbo", -elbo) logging_hook = tf.train.LoggingTensorHook( { "loss": loss, "reconstruction_loss": reconstruction_loss, "elbo": -elbo }, every_n_iter=100) return contrib_tpu.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op, training_hooks=[logging_hook]) elif mode == tf.estimator.ModeKeys.EVAL: return contrib_tpu.TPUEstimatorSpec( mode=mode, loss=loss, eval_metrics=(make_metric_fn("reconstruction_loss", "elbo", "regularizer", "kl_loss"), [ reconstruction_loss, -elbo, regularizer, kl_loss ])) else: raise NotImplementedError("Eval mode not supported.")
def test_vae_optimizer(self, gin_bindings, expected_learning_rate): gin.parse_config_files_and_bindings([], gin_bindings) with self.test_session(): x = tf.Variable(0.0) y = tf.pow(x + 2.0, 2.0) global_step = tf.train.get_or_create_global_step() optimizer = optimizers.make_vae_optimizer() train_op = optimizer.minimize(loss=y, global_step=global_step) tf.global_variables_initializer().run() for it in range(10): self.evaluate(train_op) self.assertEqual(it + 1, self.evaluate(global_step)) current_learning_rate = self.evaluate(optimizer._learning_rate_tensor) self.assertAlmostEqual(expected_learning_rate, current_learning_rate) gin.clear_config()
def model_fn(self, features, labels, mode, params): """TPUEstimator compatible model function.""" del labels is_training = (mode == tf.estimator.ModeKeys.TRAIN) data_shape = features.get_shape().as_list()[1:] z_mean, z_logvar = self.gaussian_encoder(features, is_training=is_training) z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar) z_shuffle = shuffle_codes(z_sampled) with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE): logits_z, probs_z = architectures.make_discriminator( z_sampled, is_training=is_training) _, probs_z_shuffle = architectures.make_discriminator( z_shuffle, is_training=is_training) reconstructions = self.decode(z_sampled, data_shape, is_training) per_sample_loss = losses.make_reconstruction_loss( features, reconstructions) reconstruction_loss = tf.reduce_mean(per_sample_loss) kl_loss = compute_gaussian_kl(z_mean, z_logvar) standard_vae_loss = tf.add(reconstruction_loss, kl_loss, name="VAE_loss") # tc = E[log(p_real)-log(p_fake)] = E[logit_real - logit_fake] tc_loss_per_sample = logits_z[:, 0] - logits_z[:, 1] tc_loss = tf.reduce_mean(tc_loss_per_sample, axis=0) regularizer = kl_loss + self.gamma * tc_loss factor_vae_loss = tf.add(standard_vae_loss, self.gamma * tc_loss, name="factor_VAE_loss") discr_loss = tf.add(0.5 * tf.reduce_mean(tf.log(probs_z[:, 0])), 0.5 * tf.reduce_mean(tf.log(probs_z_shuffle[:, 1])), name="discriminator_loss") if mode == tf.estimator.ModeKeys.TRAIN: optimizer_vae = optimizers.make_vae_optimizer() optimizer_discriminator = optimizers.make_discriminator_optimizer() all_variables = tf.trainable_variables() encoder_vars = [ var for var in all_variables if "encoder" in var.name ] decoder_vars = [ var for var in all_variables if "decoder" in var.name ] discriminator_vars = [var for var in all_variables \ if "discriminator" in var.name] update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) train_op_vae = optimizer_vae.minimize( loss=factor_vae_loss, global_step=tf.train.get_global_step(), var_list=encoder_vars + decoder_vars) train_op_discr = optimizer_discriminator.minimize( loss=-discr_loss, global_step=tf.train.get_global_step(), var_list=discriminator_vars) train_op = tf.group(train_op_vae, train_op_discr, update_ops) tf.summary.scalar("reconstruction_loss", reconstruction_loss) logging_hook = tf.train.LoggingTensorHook( { "loss": factor_vae_loss, "reconstruction_loss": reconstruction_loss }, every_n_iter=50) return tf.contrib.tpu.TPUEstimatorSpec( mode=mode, loss=factor_vae_loss, train_op=train_op, training_hooks=[logging_hook]) elif mode == tf.estimator.ModeKeys.EVAL: return tf.contrib.tpu.TPUEstimatorSpec( mode=mode, loss=factor_vae_loss, eval_metrics=(make_metric_fn("reconstruction_loss", "regularizer", "kl_loss"), [reconstruction_loss, regularizer, kl_loss])) else: raise NotImplementedError("Eval mode not supported.")
def model_fn(self, features, labels, mode, params): """TPUEstimator compatible model function.""" del labels is_training = (mode == tf.estimator.ModeKeys.TRAIN) data_shape = features.get_shape().as_list()[1:] output = self.gaussian_encoder(features, is_training=is_training) if len(output) == 2: z_mean, z_logvar = output else: z_mean, z_logvar, L0_reg, mask = output z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar) reconstructions = self.decode(z_sampled, data_shape, is_training) per_sample_loss = losses.make_reconstruction_loss( features, reconstructions) reconstruction_loss = tf.reduce_mean(per_sample_loss) kl_loss = compute_gaussian_kl(z_mean, z_logvar, mask) # regularizer = self.regularizer(kl_loss, z_mean, z_logvar, z_sampled, mask, L0_reg) regularizer = self.regularizer(kl_loss, z_mean, z_logvar, z_sampled) if len(output) == 2: loss = tf.add(reconstruction_loss, regularizer, name="loss") else: loss = tf.add(reconstruction_loss, regularizer + L0_reg / 500000., name="loss") loss = tf.add(reconstruction_loss, regularizer, name="loss") elbo = tf.add(reconstruction_loss, kl_loss, name="elbo") if mode == tf.estimator.ModeKeys.TRAIN: optimizer = optimizers.make_vae_optimizer() update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) train_op = optimizer.minimize( loss=loss, global_step=tf.train.get_global_step()) train_op = tf.group([train_op, update_ops]) # tf.summary.scalar("L0_reg", L0_reg) # tf.summary.scalar("mask_sum",tf.reduce_sum(mask)) tf.summary.scalar("reconstruction_loss", reconstruction_loss) tf.summary.scalar("elbo", elbo) logging_hook = tf.train.LoggingTensorHook( { "loss": loss, "reconstruction_loss": reconstruction_loss, "elbo": -elbo }, every_n_iter=100) return tf.contrib.tpu.TPUEstimatorSpec( mode=mode, loss=loss, train_op=train_op, training_hooks=[logging_hook]) elif mode == tf.estimator.ModeKeys.EVAL: return tf.contrib.tpu.TPUEstimatorSpec( mode=mode, loss=loss, eval_metrics=(make_metric_fn("reconstruction_loss", "elbo", "regularizer", "kl_loss"), [ reconstruction_loss, -elbo, regularizer, kl_loss ])) else: raise NotImplementedError("Eval mode not supported.")
def model_fn(self, features, labels, mode, params): """TPUEstimator compatible model function.""" is_training = (mode == tf.estimator.ModeKeys.TRAIN) data_shape = features.get_shape().as_list()[1:] data_shape[0] = int(data_shape[0] / 2) features_1 = features[:, :data_shape[0], :, :] features_2 = features[:, data_shape[0]:, :, :] with tf.variable_scope( tf.get_variable_scope(), reuse=tf.AUTO_REUSE): z_mean, z_logvar = self.gaussian_encoder(features_1, is_training=is_training) z_mean_2, z_logvar_2 = self.gaussian_encoder(features_2, is_training=is_training) labels = tf.squeeze(tf.one_hot(labels, z_mean.get_shape().as_list()[1])) kl_per_point = compute_kl(z_mean, z_mean_2, z_logvar, z_logvar_2) new_mean = 0.5 * z_mean + 0.5 * z_mean_2 var_1 = tf.exp(z_logvar) var_2 = tf.exp(z_logvar_2) new_log_var = tf.math.log(0.5*var_1 + 0.5*var_2) mean_sample_1, log_var_sample_1 = self.aggregate( z_mean, z_logvar, new_mean, new_log_var, labels, kl_per_point) mean_sample_2, log_var_sample_2 = self.aggregate( z_mean_2, z_logvar_2, new_mean, new_log_var, labels, kl_per_point) z_sampled_1 = self.sample_from_latent_distribution( mean_sample_1, log_var_sample_1) z_sampled_2 = self.sample_from_latent_distribution( mean_sample_2, log_var_sample_2) with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE): reconstructions_1 = self.decode(z_sampled_1, data_shape, is_training) reconstructions_2 = self.decode(z_sampled_2, data_shape, is_training) per_sample_loss_1 = losses.make_reconstruction_loss( features_1, reconstructions_1) per_sample_loss_2 = losses.make_reconstruction_loss( features_2, reconstructions_2) reconstruction_loss_1 = tf.reduce_mean(per_sample_loss_1) reconstruction_loss_2 = tf.reduce_mean(per_sample_loss_2) reconstruction_loss = (0.5 * reconstruction_loss_1 + 0.5 * reconstruction_loss_2) kl_loss_1 = vae.compute_gaussian_kl(mean_sample_1, log_var_sample_1) kl_loss_2 = vae.compute_gaussian_kl(mean_sample_2, log_var_sample_2) kl_loss = 0.5 * kl_loss_1 + 0.5 * kl_loss_2 regularizer = self.regularizer( kl_loss, None, None, None) loss = tf.add(reconstruction_loss, regularizer, name="loss") elbo = tf.add(reconstruction_loss, kl_loss, name="elbo") if mode == tf.estimator.ModeKeys.TRAIN: optimizer = optimizers.make_vae_optimizer() update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) train_op = optimizer.minimize( loss=loss, global_step=tf.train.get_global_step()) train_op = tf.group([train_op, update_ops]) tf.summary.scalar("reconstruction_loss", reconstruction_loss) tf.summary.scalar("elbo", -elbo) logging_hook = tf.train.LoggingTensorHook({ "loss": loss, "reconstruction_loss": reconstruction_loss, "elbo": -elbo, }, every_n_iter=100) return TPUEstimatorSpec( mode=mode, loss=loss, train_op=train_op, training_hooks=[logging_hook]) elif mode == tf.estimator.ModeKeys.EVAL: return TPUEstimatorSpec( mode=mode, loss=loss, eval_metrics=(make_metric_fn("reconstruction_loss", "elbo", "regularizer", "kl_loss"), [reconstruction_loss, -elbo, regularizer, kl_loss])) else: raise NotImplementedError("Eval mode not supported.")
def model_fn(self, features, labels, mode, params): """TPUEstimator compatible model function.""" is_training = (mode == tf.estimator.ModeKeys.TRAIN) labelled_features = labels[0] labels = tf.to_float(labels[1]) data_shape = features.get_shape().as_list()[1:] with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE): z_mean, z_logvar = self.gaussian_encoder(features, is_training=is_training) z_mean_labelled, _ = self.gaussian_encoder(labelled_features, is_training=is_training) supervised_loss = [] mine_ops = [] for l in range(labels.get_shape().as_list()[1]): for r in range(z_mean.get_shape().as_list()[1]): label_for_mi = tf.layers.flatten(labels[:, l]) representation_for_mi = tf.layers.flatten(z_mean_labelled[:, r]) mi_lr, op_lr = mine(representation_for_mi, label_for_mi, "estimator_network_%d_%d" % (l, r)) if l != r: supervised_loss = supervised_loss + [tf.math.square(mi_lr)] mine_ops = mine_ops + [op_lr] supervised_loss = tf.reshape(tf.add_n(supervised_loss), []) z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar) reconstructions = self.decode(z_sampled, data_shape, is_training) per_sample_loss = losses.make_reconstruction_loss( features, reconstructions) reconstruction_loss = tf.reduce_mean(per_sample_loss) kl_loss = compute_gaussian_kl(z_mean, z_logvar) standard_vae_loss = tf.add(reconstruction_loss, self.beta * kl_loss, name="VAE_loss") gamma_annealed = make_annealer(self.gamma_sup, tf.train.get_global_step()) s2_mine_vae_loss = tf.add(standard_vae_loss, gamma_annealed * supervised_loss, name="s2_factor_VAE_loss") if mode == tf.estimator.ModeKeys.TRAIN: optimizer_vae = optimizers.make_vae_optimizer() all_variables = tf.trainable_variables() encoder_vars = [ var for var in all_variables if "encoder" in var.name ] decoder_vars = [ var for var in all_variables if "decoder" in var.name ] update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) train_op_vae = optimizer_vae.minimize( loss=s2_mine_vae_loss, global_step=tf.train.get_global_step(), var_list=encoder_vars + decoder_vars) train_op = tf.group(train_op_vae, mine_ops, update_ops) tf.summary.scalar("reconstruction_loss", reconstruction_loss) logging_hook = tf.train.LoggingTensorHook( { "loss": s2_mine_vae_loss, "reconstruction_loss": reconstruction_loss, "supervised_loss": supervised_loss, }, every_n_iter=50) return TPUEstimatorSpec(mode=mode, loss=s2_mine_vae_loss, train_op=train_op, training_hooks=[logging_hook]) elif mode == tf.estimator.ModeKeys.EVAL: return TPUEstimatorSpec( mode=mode, loss=s2_mine_vae_loss, eval_metrics=(make_metric_fn("reconstruction_loss", "supervised_loss", "kl_loss"), [reconstruction_loss, supervised_loss, kl_loss])) else: raise NotImplementedError("Eval mode not supported.")
def model_fn(self, features, labels, mode, params): """TPUEstimator compatible model function. Args: features: Batch of images [batch_size, 64, 64, 3]. labels: Tuple with batch of features [batch_size, 64, 64, 3] and the labels [batch_size, labels_size]. mode: Mode for the TPUEstimator. params: Dict with parameters. Returns: TPU estimator. """ is_training = (mode == tf.estimator.ModeKeys.TRAIN) labelled_features = labels[0] labels = tf.to_float(labels[1]) data_shape = features.get_shape().as_list()[1:] with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE): z_mean, z_logvar = self.gaussian_encoder(features, is_training=is_training) z_mean_labelled, _ = self.gaussian_encoder(labelled_features, is_training=is_training) z_sampled = self.sample_from_latent_distribution(z_mean, z_logvar) reconstructions = self.decode(z_sampled, data_shape, is_training) per_sample_loss = losses.make_reconstruction_loss( features, reconstructions) reconstruction_loss = tf.reduce_mean(per_sample_loss) kl_loss = compute_gaussian_kl(z_mean, z_logvar) gamma_annealed = make_annealer(self.gamma_sup, tf.train.get_global_step()) supervised_loss = make_supervised_loss(z_mean_labelled, labels, self.factor_sizes) regularizer = self.unsupervised_regularizer( kl_loss, z_mean, z_logvar, z_sampled) + gamma_annealed * supervised_loss loss = tf.add(reconstruction_loss, regularizer, name="loss") elbo = tf.add(reconstruction_loss, kl_loss, name="elbo") if mode == tf.estimator.ModeKeys.TRAIN: optimizer = optimizers.make_vae_optimizer() update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) train_op = optimizer.minimize( loss=loss, global_step=tf.train.get_global_step()) train_op = tf.group([train_op, update_ops]) tf.summary.scalar("reconstruction_loss", reconstruction_loss) tf.summary.scalar("elbo", -elbo) logging_hook = tf.train.LoggingTensorHook( { "loss": loss, "reconstruction_loss": reconstruction_loss, "elbo": -elbo, "supervised_loss": supervised_loss }, every_n_iter=100) return TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op, training_hooks=[logging_hook]) elif mode == tf.estimator.ModeKeys.EVAL: return TPUEstimatorSpec(mode=mode, loss=loss, eval_metrics=(make_metric_fn( "reconstruction_loss", "elbo", "regularizer", "kl_loss", "supervised_loss"), [ reconstruction_loss, -elbo, regularizer, kl_loss, supervised_loss ])) else: raise NotImplementedError("Eval mode not supported.")