Adding flag to set random seeds.

Victor Bittorf · Victor Bittorf · commit babbb38f037e · 2018-04-11T15:19:28.000-07:00
diff --git a/official/mnist/mnist.py b/official/mnist/mnist.py
@@ -182,6 +182,9 @@ def main(argv):
 
   model_function = model_fn
 
+  if flags.seed is not None:
+    model_helpers.set_random_seed(flags.seed)
+
   if flags.multi_gpu:
     validate_batch_size_for_multi_gpu(flags.batch_size)
 
diff --git a/official/resnet/imagenet_main.py b/official/resnet/imagenet_main.py
@@ -26,6 +26,7 @@
 from official.resnet import imagenet_preprocessing
 from official.resnet import resnet_model
 from official.resnet import resnet_run_loop
+from official.utils.misc import model_helpers
 
 _DEFAULT_IMAGE_SIZE = 224
 _NUM_CHANNELS = 3
@@ -315,6 +316,9 @@ def main(argv):
 
   flags = parser.parse_args(args=argv[1:])
 
+  if flags.seed is not None:
+    model_helpers.set_random_seed(flags.seed)
+
   input_function = flags.use_synthetic_data and get_synth_input_fn() or input_fn
 
   resnet_run_loop.resnet_main(
diff --git a/official/utils/arg_parsers/parsers.py b/official/utils/arg_parsers/parsers.py
@@ -104,12 +104,13 @@ class BaseParser(argparse.ArgumentParser):
     batch_size: Create a flag to specify the batch size.
     multi_gpu: Create a flag to allow the use of all available GPUs.
     hooks: Create a flag to specify hooks for logging.
+    seed: Create a flag to set random seeds.
   """
 
   def __init__(self, add_help=False, data_dir=True, model_dir=True,
                train_epochs=True, epochs_between_evals=True,
                stop_threshold=True, batch_size=True, multi_gpu=True,
-               hooks=True):
+               hooks=True, seed=True):
     super(BaseParser, self).__init__(add_help=add_help)
 
     if data_dir:
@@ -176,6 +177,16 @@ def __init__(self, add_help=False, data_dir=True, model_dir=True,
           metavar="<HK>"
       )
 
+    if seed:
+      self.add_argument(
+          "--seed", "-s", nargs="+", type=int, default=None,
+          help="[default: %(default)s] An integer to seed random number"
+               "generators. If unset, RNGs choose their own seeds resulting "
+               "in each run having a different seed.",
+          metavar="<SEED>"
+      )
+
+
 
 class PerformanceParser(argparse.ArgumentParser):
   """Default parser for specifying performance tuning arguments.
diff --git a/official/utils/misc/model_helpers.py b/official/utils/misc/model_helpers.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import numbers
+import random
 
 import tensorflow as tf
 
@@ -53,3 +54,25 @@ def past_stop_threshold(stop_threshold, eval_metric):
     return True
 
   return False
+
+
+def set_random_seed(seed):
+  """Sets the random seeds for available RNGs.
+  This seeds RNGs for python's random and for Tensorflow. The intended
+  use case is for this to be called exactly once at the start of execution
+  to improve stability and reproducability between runs.
+
+  Successive calls to re-seed will not behave as expected. This should
+  be called at most once.
+
+  Args:
+    seed: integer, a seed which will be passed to the RNGs.
+
+  Raises:
+    ValueError: if the seed is not an integer or if deemed unsuitable for
+    seeding a the RNGs.
+  """
+  if not isinstance(seed, int):
+    raise ValueError("Random seed is not an integer: {}".format(seed))
+  random.seed(seed)
+  tf.set_random_seed(seed)
diff --git a/official/utils/misc/model_helpers_test.py b/official/utils/misc/model_helpers_test.py
@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import random
+
 import tensorflow as tf  # pylint: disable=g-bad-import-order
 
 from official.utils.misc import model_helpers
@@ -64,6 +66,29 @@ def test_past_stop_threshold_not_number(self):
     with self.assertRaises(ValueError):
       model_helpers.past_stop_threshold(tf.constant(4), None)
 
+  def test_random_seed(self):
+    """It is unclear if this test is a good idea or stable.
+    If tests are run in parallel, this could be flakey."""
+    model_helpers.set_random_seed(42)
+    expected_py_random = [int(random.random() * 1000) for i in range(10)]
+    tf_random = []
+    with tf.Session() as sess:
+      for i in range(10):
+          a = tf.random_uniform([1])
+          tf_random.append(int(sess.run(a)[0] * 1000))
+
+    model_helpers.set_random_seed(42)
+    py_random = [int(random.random() * 1000) for i in range(10)]
+
+    # Instead of concerning ourselves with the particular results, we simply
+    # want to ensure that the results are reproducible. So, we seed, read,
+    # re-seed, re-read.
+    self.assertAllEqual(expected_py_random, py_random)
+
+    # TF does not accept being re-seeded.
+    expected_tf_random = [637, 689, 961, 969, 321, 390, 919, 681, 112, 187]
+    self.assertAllEqual(expected_tf_random, tf_random)
+
 
 if __name__ == "__main__":
   tf.test.main()
diff --git a/official/wide_deep/wide_deep.py b/official/wide_deep/wide_deep.py
@@ -179,6 +179,9 @@ def main(argv):
   parser = WideDeepArgParser()
   flags = parser.parse_args(args=argv[1:])
 
+  if flags.seed is not None:
+    model_helpers.set_random_seed(flags.seed)
+
   # Clean up the model directory if present
   shutil.rmtree(flags.model_dir, ignore_errors=True)
   model = build_estimator(flags.model_dir, flags.model_type)
