def testTransformerAutoencoder(self):
hparams = imagetransformer_latent_tiny()
hparams.mode = tf.estimator.ModeKeys.TRAIN
block_dim = int(hparams.hidden_size // hparams.num_blocks)
block_v_size = 2**(hparams.bottleneck_bits /
(hparams.num_residuals * hparams.num_blocks))
block_v_size = int(block_v_size)
means = tf.get_variable(
name="means",
shape=[hparams.num_residuals,
hparams.num_blocks,
block_v_size,
block_dim],
initializer=tf.uniform_unit_scaling_initializer())
hparams.bottleneck = functools.partial(
discretization.discrete_bottleneck,
hidden_size=hparams.hidden_size,
z_size=hparams.bottleneck_bits,
filter_size=hparams.filter_size,
startup_steps=hparams.startup_steps,
bottleneck_kind=hparams.bottleneck_kind,
num_blocks=hparams.num_blocks,
num_residuals=hparams.num_residuals,
reshape_method=hparams.reshape_method,
beta=hparams.vq_beta,
decay=hparams.vq_decay,
soft_em=hparams.soft_em,
num_samples=hparams.num_samples,
epsilon=hparams.vq_epsilon,
ema=hparams.ema,
means=means)
inputs = None
batch_size = hparams.batch_size
targets = tf.random_uniform([batch_size,
hparams.img_len,
hparams.img_len,
hparams.hidden_size],
minval=-1., maxval=1.)
target_space_id = None
tf.train.create_global_step()
decoder_output, losses, cache = latent_layers.transformer_autoencoder(
inputs, targets, target_space_id, hparams)
self.assertEqual(set(six.iterkeys(losses)),
{"extra", "extra_loss", "latent_pred"})
self.evaluate(tf.global_variables_initializer())
decoder_output_, extra_loss_, latent_pred_ = self.evaluate(
[decoder_output, losses["extra_loss"], losses["latent_pred"]])
self.assertEqual(decoder_output_.shape, (batch_size,
hparams.img_len,
hparams.img_len,
hparams.hidden_size))
self.assertEqual(extra_loss_.shape, (batch_size,))
self.assertEqual(latent_pred_.shape, (batch_size,))
self.assertAllGreaterEqual(extra_loss_, 0.)
self.assertAllGreaterEqual(latent_pred_, 0.)
self.assertEqual(cache, None)
def testTransformerAutoencoder(self):
hparams = imagetransformer_latent_tiny()
hparams.mode = tf.estimator.ModeKeys.TRAIN
block_dim = int(hparams.hidden_size // hparams.num_blocks)
block_v_size = 2**(hparams.bottleneck_bits /
(hparams.num_residuals * hparams.num_blocks))
block_v_size = int(block_v_size)
means = tf.get_variable(
name="means",
shape=[hparams.num_residuals,
hparams.num_blocks,
block_v_size,
block_dim],
initializer=tf.uniform_unit_scaling_initializer())
hparams.bottleneck = functools.partial(
discretization.discrete_bottleneck,
hidden_size=hparams.hidden_size,
z_size=hparams.bottleneck_bits,
filter_size=hparams.filter_size,
startup_steps=hparams.startup_steps,
bottleneck_kind=hparams.bottleneck_kind,
num_blocks=hparams.num_blocks,
num_residuals=hparams.num_residuals,
reshape_method=hparams.reshape_method,
beta=hparams.vq_beta,
decay=hparams.vq_decay,
soft_em=hparams.soft_em,
num_samples=hparams.num_samples,
epsilon=hparams.vq_epsilon,
ema=hparams.ema,
means=means)
inputs = None
batch_size = hparams.batch_size
targets = tf.random_uniform([batch_size,
hparams.img_len,
hparams.img_len,
hparams.hidden_size],
minval=-1., maxval=1.)
target_space_id = None
tf.train.create_global_step()
decoder_output, losses, cache = latent_layers.transformer_autoencoder(
inputs, targets, target_space_id, hparams)
self.assertEqual(set(six.iterkeys(losses)),
{"extra", "extra_loss", "latent_pred"})
self.evaluate(tf.global_variables_initializer())
decoder_output_, extra_loss_, latent_pred_ = self.evaluate(
[decoder_output, losses["extra_loss"], losses["latent_pred"]])
self.assertEqual(decoder_output_.shape, (batch_size,
hparams.img_len,
hparams.img_len,
hparams.hidden_size))
self.assertEqual(extra_loss_.shape, (batch_size,))
self.assertEqual(latent_pred_.shape, (batch_size,))
self.assertAllGreaterEqual(extra_loss_, 0.)
self.assertAllGreaterEqual(latent_pred_, 0.)
self.assertEqual(cache, None)