Add random seed, py3 fix, disable flaky test

PiperOrigin-RevId: 179942374

Author: Ryan Sepassi

.travis.yml

@@ -14,7 +14,7 @@ env:
     - T2T_DATA_DIR=/tmp/t2t-data
     - T2T_TRAIN_DIR=/tmp/t2t-train
 script:
-  - pytest --ignore=tensor2tensor/utils/registry_test.py --ignore=tensor2tensor/problems_test.py --ignore=tensor2tensor/tpu/tpu_trainer_lib_test.py
+  - pytest --ignore=tensor2tensor/utils/registry_test.py --ignore=tensor2tensor/problems_test.py --ignore=tensor2tensor/tpu/tpu_trainer_lib_test.py --ignore=tensor2tensor/data_generators/algorithmic_math_test.py
   - pytest tensor2tensor/utils/registry_test.py
   - pytest tensor2tensor/tpu/tpu_trainer_lib_test.py
   - t2t-datagen 2>&1 | grep translate && echo passed

tensor2tensor/bin/t2t-trainer

@@ -45,6 +45,7 @@ flags.DEFINE_string("t2t_usr_dir", "",
                     "The imported files should contain registrations, "
                     "e.g. @registry.register_model calls, that will then be "
                     "available to the t2t-trainer.")
+flags.DEFINE_integer("random_seed", 1234, "Random seed.")
 flags.DEFINE_integer("tpu_num_shards", 8, "Number of tpu shards.")
 flags.DEFINE_integer("iterations_per_loop", 1000,
                      "Number of iterations in a TPU training loop.")
@@ -171,7 +172,7 @@ def execute_schedule(exp):
 
 def main(_):
   tf.logging.set_verbosity(tf.logging.INFO)
-  tf.set_random_seed(123)
+  tpu_trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   log_registry()
 

tensor2tensor/bin/t2t_trainer.py

@@ -44,6 +44,7 @@
                     "The imported files should contain registrations, "
                     "e.g. @registry.register_model calls, that will then be "
                     "available to the t2t-trainer.")
+flags.DEFINE_integer("random_seed", 1234, "Random seed.")
 flags.DEFINE_integer("tpu_num_shards", 8, "Number of tpu shards.")
 flags.DEFINE_integer("iterations_per_loop", 1000,
                      "Number of iterations in a TPU training loop.")
@@ -170,7 +171,7 @@ def execute_schedule(exp):
 
 def main(_):
   tf.logging.set_verbosity(tf.logging.INFO)
-  tf.set_random_seed(123)
+  tpu_trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   log_registry()
 

tensor2tensor/data_generators/algorithmic_math_test.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 """Tests for tensor2tensor.data_generators.algorithmic_math."""
+# TODO(rsepassi): This test is flaky. Disable, remove, or update.
 
 from __future__ import absolute_import
 from __future__ import division

tensor2tensor/data_generators/problem.py

@@ -947,7 +947,7 @@ def standardize_shapes(features, batch_size=None):
 
 def pad_batch(features, batch_multiple):
   """Pad batch dim of features to nearest multiple of batch_multiple."""
-  feature = features.items()[0][1]
+  feature = list(features.items())[0][1]
   batch_size = tf.shape(feature)[0]
   mod = batch_size % batch_multiple
   has_mod = tf.cast(tf.cast(mod, tf.bool), tf.int32)

tensor2tensor/tpu/tpu_trainer.py

@@ -44,6 +44,7 @@
                     "The imported files should contain registrations, "
                     "e.g. @registry.register_model calls, that will then be "
                     "available to the t2t-trainer.")
+flags.DEFINE_integer("random_seed", 1234, "Random seed.")
 flags.DEFINE_integer("tpu_num_shards", 8, "Number of tpu shards.")
 flags.DEFINE_integer("iterations_per_loop", 1000,
                      "Number of iterations in a TPU training loop.")
@@ -170,7 +171,7 @@ def execute_schedule(exp):
 
 def main(_):
   tf.logging.set_verbosity(tf.logging.INFO)
-  tf.set_random_seed(123)
+  tpu_trainer_lib.set_random_seed(FLAGS.random_seed)
   usr_dir.import_usr_dir(FLAGS.t2t_usr_dir)
   log_registry()
 

tensor2tensor/tpu/tpu_trainer_lib.py

@@ -20,9 +20,12 @@
 from __future__ import print_function
 
 import os
+import random
 
 # Dependency imports
 
+import numpy as np
+
 from tensor2tensor.utils import devices
 from tensor2tensor.utils import expert_utils
 from tensor2tensor.utils import metrics_hook
@@ -336,3 +339,9 @@ def add_problem_hparams(hparams, problems):
 
     hparams.problem_instances.append(problem)
     hparams.problems.append(p_hparams)
+
+
+def set_random_seed(seed):
+  tf.set_random_seed(seed)
+  random.seed(seed)
+  np.random.seed(seed)