Add t2t_usr_dir functionality to t2t-datagen

PiperOrigin-RevId: 162253918

Author: Ryan Sepassi

tensor2tensor/bin/t2t-datagen

@@ -48,6 +48,7 @@ from tensor2tensor.data_generators import wiki
 from tensor2tensor.data_generators import wmt
 from tensor2tensor.data_generators import wsj_parsing
 from tensor2tensor.utils import registry
+from tensor2tensor.utils import registry_utils
 
 import tensorflow as tf
 
@@ -63,6 +64,12 @@ flags.DEFINE_integer("num_shards", 10, "How many shards to use.")
 flags.DEFINE_integer("max_cases", 0,
                      "Maximum number of cases to generate (unbounded if 0).")
 flags.DEFINE_integer("random_seed", 429459, "Random seed to use.")
+flags.DEFINE_string("t2t_usr_dir", "",
+                    "Path to a Python module that will be imported. The "
+                    "__init__.py file should include the necessary imports. "
+                    "The imported files should contain registrations, "
+                    "e.g. @registry.register_problem calls, that will then be "
+                    "available to t2t-datagen.")
 
 # Mapping from problems that we can generate data for to their generators.
 # pylint: disable=g-long-lambda
@@ -273,6 +280,7 @@ def set_random_seed():
 
 def main(_):
   tf.logging.set_verbosity(tf.logging.INFO)
+  registry_utils.import_usr_dir(FLAGS.t2t_usr_dir)
 
   # Calculate the list of problems to generate.
   problems = sorted(

tensor2tensor/bin/t2t-trainer

@@ -29,12 +29,9 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import importlib
-import os
-import sys
-
 # Dependency imports
 
+from tensor2tensor.utils import registry_utils
 from tensor2tensor.utils import trainer_utils as utils
 
 import tensorflow as tf
@@ -50,24 +47,9 @@ flags.DEFINE_string("t2t_usr_dir", "",
                     "available to the t2t-trainer.")
 
 
-def import_usr_dir():
-  """Import module at FLAGS.t2t_usr_dir, if provided."""
-  if not FLAGS.t2t_usr_dir:
-    return
-  dir_path = os.path.expanduser(FLAGS.t2t_usr_dir)
-  if dir_path[-1] == "/":
-    dir_path = dir_path[:-1]
-  containing_dir, module_name = os.path.split(dir_path)
-  tf.logging.info("Importing user module %s from path %s", module_name,
-                  containing_dir)
-  sys.path.insert(0, containing_dir)
-  importlib.import_module(module_name)
-  sys.path.pop(0)
-
-
 def main(_):
   tf.logging.set_verbosity(tf.logging.INFO)
-  import_usr_dir()
+  registry_utils.import_usr_dir(FLAGS.t2t_usr_dir)
   utils.log_registry()
   utils.validate_flags()
   utils.run(

tensor2tensor/utils/registry_utils.py

@@ -0,0 +1,41 @@
+# Copyright 2017 The Tensor2Tensor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for the t2t registry."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import importlib
+import os
+import sys
+
+# Dependency imports
+
+import tensorflow as tf
+
+
+def import_usr_dir(usr_dir):
+  """Import module at usr_dir, if provided."""
+  if not usr_dir:
+    return
+  dir_path = os.path.expanduser(usr_dir)
+  if dir_path[-1] == "/":
+    dir_path = dir_path[:-1]
+  containing_dir, module_name = os.path.split(dir_path)
+  tf.logging.info("Importing user module %s from path %s", module_name,
+                  containing_dir)
+  sys.path.insert(0, containing_dir)
+  importlib.import_module(module_name)
+  sys.path.pop(0)