Fix for issue #215 on github, update transformer_vae.

PiperOrigin-RevId: 164771762

Author: Lukasz Kaiser

tensor2tensor/models/transformer_vae.py

@@ -109,7 +109,7 @@ def encode(x, x_space, hparams, name):
   with tf.variable_scope(name):
     (encoder_input, encoder_self_attention_bias,
      _) = transformer.transformer_prepare_encoder(x, x_space, hparams)
-    encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.residual_dropout)
+    encoder_input = tf.nn.dropout(encoder_input, 1.0 - hparams.dropout)
     return transformer.transformer_encoder(
         encoder_input, encoder_self_attention_bias, hparams)
 
@@ -143,7 +143,7 @@ def vae_transformer_internal(inputs, targets, target_space, hparams):
     max_prestep = hparams.kl_warmup_steps
     prob_targets = 0.95 if is_training else 1.0
     targets_dropout_max = common_layers.inverse_lin_decay(max_prestep) - 0.01
-    targets = dropmask(targets, targets_dropout_max, is_training)
+    targets = dropmask(targets, targets_dropout_max * 0.7, is_training)
     targets = tf.cond(tf.less(tf.random_uniform([]), prob_targets),
                       lambda: targets, lambda: tf.zeros_like(targets))
 
@@ -168,7 +168,7 @@ def vae_transformer_internal(inputs, targets, target_space, hparams):
     # ret = tf.squeeze(to_decode, axis=2)
 
     # Randomize decoder inputs..
-    kl_loss *= common_layers.inverse_exp_decay(max_prestep) * 3.0
+    kl_loss *= common_layers.inverse_exp_decay(max_prestep) * 10.0
     return tf.expand_dims(ret, axis=2), kl_loss
 
 

tensor2tensor/utils/devices.py

@@ -112,7 +112,7 @@ def _replica_device_setter(worker_device):
   if FLAGS.schedule == "local_run":
     assert not FLAGS.sync
     datashard_devices = ["gpu:%d" % d for d in _gpu_order(FLAGS.worker_gpu)]
-    if FLAGS.locally_shard_to_cpu:
+    if FLAGS.locally_shard_to_cpu or FLAGS.worker_gpu < 1:
       datashard_devices += ["cpu:0"]
     caching_devices = None
   elif FLAGS.sync: