TF Eager improvements for T2TModel

PiperOrigin-RevId: 177641254

Author: Ryan Sepassi

tensor2tensor/layers/common_hparams.py

@@ -184,9 +184,6 @@ def basic_params1():
       # This is the actual batch size, *not* tokens per batch (i.e. for
       # language models this is the number of sentences in the batch)
       tpu_batch_size_per_shard=24,
-      # Things not compatible with eager mode use this flag to implement
-      # alternative functionality. We expect this to go away soon.
-      use_eager_mode=False,
       # Set by tpu_trainer to let the model know whether we are on TPU.
       # Switching on/off tpu should not invalidate checkpoints.
       use_tpu=False,

tensor2tensor/layers/common_layers.py

@@ -32,6 +32,7 @@
 
 import tensorflow as tf
 
+from tensorflow.python.eager import context as tfe_context
 from tensorflow.python.framework import function
 from tensorflow.python.framework import ops
 
@@ -200,16 +201,15 @@ def flatten4d3d(x):
   return result
 
 
-def embedding(x, vocab_size, dense_size, name=None, reuse=None, multiplier=1.0,
-              use_eager_mode=False):
+def embedding(x, vocab_size, dense_size, name=None, reuse=None, multiplier=1.0):
   """Embed x of type int64 into dense vectors, reducing to max 4 dimensions."""
   with tf.variable_scope(
       name, default_name="embedding", values=[x], reuse=reuse):
     embedding_var = tf.get_variable("kernel", [vocab_size, dense_size])
     # On the backwards pass, we want to convert the gradient from
     # an indexed-slices to a regular tensor before sending it back to the
     # parameter server. This avoids excess computation on the parameter server.
-    if not use_eager_mode:
+    if not tfe_context.in_eager_mode():
       embedding_var = eu.convert_gradient_to_tensor(embedding_var)
     emb_x = tf.gather(embedding_var, x)
     if multiplier != 1.0:

tensor2tensor/layers/modalities.py

@@ -29,6 +29,8 @@
 
 import tensorflow as tf
 
+from tensorflow.python.eager import context
+
 
 # TODO(noam): remove this function after TPUs do gather faster.
 def tpu_gather(params, indices):
@@ -96,7 +98,7 @@ def _get_weights(self, hidden_dim=None):
     else:
       ret = tf.concat(shards, 0)
     # Convert ret to tensor.
-    if not self._model_hparams.use_eager_mode:
+    if not context.in_eager_mode():
       ret = eu.convert_gradient_to_tensor(ret)
     return ret
 
@@ -205,7 +207,7 @@ class ImageModality(modality.Modality):
   def bottom(self, inputs):
     with tf.variable_scope(self.name):
       inputs = common_layers.standardize_images(inputs)
-      if not self._model_hparams.use_eager_mode:
+      if not context.in_eager_mode():
         tf.summary.image("inputs", inputs, max_outputs=2)
       return tf.to_float(inputs)
 
@@ -216,8 +218,7 @@ def targets_bottom(self, inputs):
           tf.to_int32(common_layers.flatten4d3d(inputs)),
           self.top_dimensionality,
           self._body_input_depth,
-          name="input_rgb_embedding",
-          use_eager_mode=self._model_hparams.use_eager_mode)
+          name="input_rgb_embedding")
       if self._model_hparams.multiply_embedding_mode == "sqrt_depth":
         ret *= self._body_input_depth**0.5
 

tensor2tensor/layers/modalities_test.py

@@ -43,7 +43,6 @@ def testSymbolModalityInputs(self):
         symbol_modality_skip_top=0,
         shared_embedding_and_softmax_weights=0,
         prepend_mode="none",
-        use_eager_mode=False,
         use_tpu=False)
     x = -1 + np.random.random_integers(
         vocab_size, size=(batch_size, length, 1, 1))
@@ -74,7 +73,6 @@ def testSymbolModalityTargets(self):
         factored_logits=0,
         mode=tf.estimator.ModeKeys.TRAIN,
         prepend_mode="none",
-        use_eager_mode=False,
         use_tpu=False)
     body_output = -1 + np.random.random_integers(
         100, size=(batch_size, length, height, hidden_size))
@@ -112,7 +110,6 @@ def testSymbolModalityTargetsFactored(self):
         factored_logits=1,
         mode=tf.estimator.ModeKeys.TRAIN,
         prepend_mode="none",
-        use_eager_mode=False,
         use_tpu=False)
     body_output = -1 + np.random.random_integers(
         100, size=(batch_size, length, height, hidden_size))

tensor2tensor/models/cycle_gan.py

@@ -66,11 +66,10 @@ def cycle_gan_internal(inputs, targets, _, hparams):
     # Embed inputs and targets.
     inputs_orig, targets_orig = tf.to_int32(inputs), tf.to_int32(targets)
     inputs = common_layers.embedding(
-        inputs_orig, hparams.vocab_size, hparams.hidden_size, "embed",
-        use_eager_mode=hparams.use_eager_mode)
+        inputs_orig, hparams.vocab_size, hparams.hidden_size, "embed")
     targets = common_layers.embedding(
         targets_orig, hparams.vocab_size, hparams.hidden_size,
-        "embed", reuse=True, use_eager_mode=hparams.use_eager_mode)
+        "embed", reuse=True)
 
     # Split the batch into input-input and target-target parts.
     inputs1, _ = split_on_batch(inputs)

tensor2tensor/models/transformer.py

@@ -37,6 +37,7 @@
 
 import tensorflow as tf
 
+from tensorflow.python.eager import context
 from tensorflow.python.util import nest
 
 
@@ -324,7 +325,7 @@ def symbols_to_logits_fn(ids, i, cache):
     # Note: Tensor.set_shape() does not work here since it merges shape info.
     # TODO(llion); Find a more robust solution.
     # pylint: disable=protected-access
-    if not self._hparams.use_eager_mode:
+    if not context.in_eager_mode():
       for layer in cache:
         cache[layer]["k"]._shape = tf.TensorShape([None, None, key_channels])
         cache[layer]["v"]._shape = tf.TensorShape([None, None, value_channels])
@@ -452,8 +453,7 @@ def transformer_prepare_encoder(inputs, target_space, hparams, features=None):
         common_layers.shape_list(inputs)[1])
   # Append target_space_id embedding to inputs.
   emb_target_space = common_layers.embedding(
-      target_space, 32, ishape_static[-1], name="target_space_embedding",
-      use_eager_mode=hparams.use_eager_mode)
+      target_space, 32, ishape_static[-1], name="target_space_embedding")
   emb_target_space = tf.reshape(emb_target_space, [1, 1, -1])
   encoder_input += emb_target_space
   if hparams.pos == "timing":