Changes needed for T2T's problems to be imported on TF 2.0 -- this do…

…esn't mean we are using these problems. This is only so that we can `import trax` on these environments. PiperOrigin-RevId: 282108974
tensorflow · Nov 23, 2019 · 8107e53 · 8107e53
1 parent 313fdfc
commit 8107e53
Show file tree

Hide file tree

Showing 22 changed files with 67 additions and 69 deletions.
diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name='tensor2tensor',
-    version='1.15.1',
+    version='1.15.2',
     description='Tensor2Tensor',
     long_description=(
         'Tensor2Tensor, or T2T for short, is a library of '
@@ -66,7 +66,6 @@
         'tensorflow-datasets',
         'tensorflow-gan',
         'tensorflow-probability==0.7.0',
-        'tf_slim',
         'tqdm',
     ],
     extras_require={

diff --git a/tensor2tensor/data_generators/allen_brain.py b/tensor2tensor/data_generators/allen_brain.py
@@ -44,7 +44,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 _BASE_EXAMPLE_IMAGE_SIZE = 64
 
@@ -351,7 +350,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "targets":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),

diff --git a/tensor2tensor/data_generators/bair_robot_pushing.py b/tensor2tensor/data_generators/bair_robot_pushing.py
@@ -36,7 +36,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 DATA_URL = (
     "http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar")
@@ -103,7 +102,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders
@@ -188,9 +187,9 @@ def extra_reading_spec(self):
         "action": tf.FixedLenFeature([4], tf.float32),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
-        "action": slim.tfexample_decoder.Tensor(tensor_key="action"),
+        "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
     }
     return data_fields, decoders
 
diff --git a/tensor2tensor/data_generators/fsns.py b/tensor2tensor/data_generators/fsns.py
@@ -28,7 +28,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 @registry.register_problem
@@ -77,5 +76,5 @@ def example_reading_spec(self):
         super(ImageFSNS, self).example_reading_spec())
     data_fields[label_key] = tf.VarLenFeature(tf.int64)
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
diff --git a/tensor2tensor/data_generators/gym_env.py b/tensor2tensor/data_generators/gym_env.py
@@ -36,7 +36,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 Frame = collections.namedtuple(
@@ -378,7 +377,7 @@ def extra_reading_spec(self):
         name: tf.FixedLenFeature([1], tf.int64) for name in field_names
     }
     decoders = {
-        name: slim.tfexample_decoder.Tensor(tensor_key=name)
+        name: tf.contrib.slim.tfexample_decoder.Tensor(tensor_key=name)
         for name in field_names
     }
     return (data_fields, decoders)

diff --git a/tensor2tensor/data_generators/image_utils.py b/tensor2tensor/data_generators/image_utils.py
@@ -31,7 +31,6 @@
 from tensor2tensor.utils import metrics
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 def matplotlib_pyplot():
@@ -173,7 +172,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -240,7 +239,7 @@ def example_reading_spec(self):
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
 
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def hparams(self, defaults, unused_model_hparams):
@@ -344,7 +343,7 @@ def example_reading_spec(self):
         super(Image2TextProblem, self).example_reading_spec())
     data_fields[label_key] = tf.VarLenFeature(tf.int64)
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def feature_encoders(self, data_dir):

diff --git a/tensor2tensor/data_generators/moving_mnist.py b/tensor2tensor/data_generators/moving_mnist.py
@@ -37,7 +37,6 @@
 import tensorflow as tf
 import tensorflow_datasets as tfds
 from tensorflow_datasets.video import moving_sequence
-import tf_slim as slim
 
 
 DATA_URL = (
@@ -95,7 +94,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders

diff --git a/tensor2tensor/data_generators/problem.py b/tensor2tensor/data_generators/problem.py
@@ -32,8 +32,13 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
-import tf_slim as slim
-from tensorflow.contrib.tpu.python.tpu import tpu_config
+# pylint: disable=g-import-not-at-top
+try:
+  from tensorflow.contrib.tpu.python.tpu import tpu_config
+except ImportError:
+  # TF 2.0 doesn't ship with contrib.
+  tpu_config = None
+# pylint: enable=g-import-not-at-top
 
 
 
@@ -199,7 +204,7 @@ class Problem(object):
         - Mutate defaults as needed
     * example_reading_spec
         - Specify the names and types of the features on disk.
-        - Specify slim.tfexample_decoder
+        - Specify tf.contrib.slim.tfexample_decoder
     * preprocess_example(example, mode, hparams)
         - Preprocess the example feature dict from feature name to Tensor or
           SparseTensor.
@@ -643,7 +648,7 @@ def dataset(self,
 
     data_filepattern = self.filepattern(data_dir, dataset_split, shard=shard)
     tf.logging.info("Reading data files from %s", data_filepattern)
-    data_files = sorted(slim.parallel_reader.get_data_files(
+    data_files = sorted(tf.contrib.slim.parallel_reader.get_data_files(
         data_filepattern))
 
     # Functions used in dataset transforms below. `filenames` can be either a
@@ -706,12 +711,12 @@ def decode_example(self, serialized_example):
     data_fields["batch_prediction_key"] = tf.FixedLenFeature([1], tf.int64, 0)
     if data_items_to_decoders is None:
       data_items_to_decoders = {
-          field: slim.tfexample_decoder.Tensor(field)
+          field: tf.contrib.slim.tfexample_decoder.Tensor(field)
           for field in data_fields
       }
 
-    decoder = slim.tfexample_decoder.TFExampleDecoder(data_fields,
-                                                      data_items_to_decoders)
+    decoder = tf.contrib.slim.tfexample_decoder.TFExampleDecoder(
+        data_fields, data_items_to_decoders)
 
     decode_items = list(sorted(data_items_to_decoders))
     decoded = decoder.decode(serialized_example, items=decode_items)

diff --git a/tensor2tensor/data_generators/style_transfer.py b/tensor2tensor/data_generators/style_transfer.py
@@ -32,16 +32,12 @@
 from tensor2tensor.data_generators import text_problems
 from tensor2tensor.utils import registry
 
-import tensorflow as tf
 
-logger = tf.logging
+# Modern-Shakespeare corpus is consisted of:
+# - 18,395 parallel sentences for training (train set),
+# - 1,218 parallel sentences for evaluation (dev set),
+# - 1,462 parallel sentence for testing (test set).
 
-"""
-Modern-Shakespeare corpus is consisted of:
-- 18,395 parallel sentences for training (train set),
-- 1,218 parallel sentences for evaluation (dev set),
-- 1,462 parallel sentence for testing (test set).
-"""
 
 _SHAKESPEARE_MODERN_TRAIN_DATASET = [[
     "https://github.com/tlatkowski/st/raw/master/shakespeare.train.tgz",

diff --git a/tensor2tensor/data_generators/translate.py b/tensor2tensor/data_generators/translate.py
@@ -32,7 +32,6 @@
 from tensor2tensor.utils import mlperf_log
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 class TranslateProblem(text_problems.Text2TextProblem):
@@ -277,8 +276,8 @@ def example_reading_spec(self):
 
     # hack: ignoring true targets and putting dist_targets in targets
     data_items_to_decoders = {
-        "inputs": slim.tfexample_decoder.Tensor("inputs"),
-        "targets": slim.tfexample_decoder.Tensor("dist_targets"),
+        "inputs": tf.contrib.slim.tfexample_decoder.Tensor("inputs"),
+        "targets": tf.contrib.slim.tfexample_decoder.Tensor("dist_targets"),
     }
 
     return (data_fields, data_items_to_decoders)

diff --git a/tensor2tensor/data_generators/video_generated.py b/tensor2tensor/data_generators/video_generated.py
@@ -28,7 +28,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 try:
   import matplotlib  # pylint: disable=g-import-not-at-top
@@ -86,7 +85,7 @@ def extra_reading_spec(self):
         "frame_number": tf.FixedLenFeature([1], tf.int64),
     }
     decoders = {
-        "frame_number": slim.tfexample_decoder.Tensor(
+        "frame_number": tf.contrib.slim.tfexample_decoder.Tensor(
             tensor_key="frame_number"),
     }
     return data_fields, decoders

diff --git a/tensor2tensor/data_generators/video_utils.py b/tensor2tensor/data_generators/video_utils.py
@@ -35,7 +35,6 @@
 from tensor2tensor.utils import metrics
 from tensor2tensor.utils import video_metrics
 import tensorflow as tf
-import tf_slim as slim
 
 
 FLAGS = flags.FLAGS
@@ -385,7 +384,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "frame":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 shape=[self.frame_height, self.frame_width, self.num_channels],
@@ -677,7 +676,7 @@ def example_reading_spec(self):
 
     data_items_to_decoders = {
         "inputs":
-            slim.tfexample_decoder.Image(
+            tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
                 channels=self.num_channels),
@@ -767,7 +766,7 @@ def example_reading_spec(self):
         super(Video2ClassProblem, self).example_reading_spec())
     data_fields[label_key] = tf.FixedLenFeature((1,), tf.int64)
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(label_key)
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(label_key)
     return data_fields, data_items_to_decoders
 
   def hparams(self, defaults, unused_model_hparams):

diff --git a/tensor2tensor/data_generators/vqa.py b/tensor2tensor/data_generators/vqa.py
@@ -39,7 +39,6 @@
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
-import tf_slim as slim
 
 
 def _get_vqa_v2_annotations(directory,
@@ -218,10 +217,10 @@ def example_reading_spec(self):
         (), tf.int64, allow_missing=True)
 
     data_items_to_decoders[
-        "question"] = slim.tfexample_decoder.Tensor(
+        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/question")
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/answer")
     return data_fields, data_items_to_decoders
 
@@ -339,23 +338,23 @@ def example_reading_spec(self):
         (), tf.int64, allow_missing=True)
 
     data_items_to_decoders[
-        "inputs"] = slim.tfexample_decoder.Tensor(
+        "inputs"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/feature")
     data_items_to_decoders[
-        "question_id"] = slim.tfexample_decoder.Tensor(
+        "question_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/question_id")
     data_items_to_decoders[
-        "image_id"] = slim.tfexample_decoder.Tensor(
+        "image_id"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/image_id")
 
     data_items_to_decoders[
-        "spatial_feature"] = slim.tfexample_decoder.Tensor(
+        "spatial_feature"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/spatial_feature")
     data_items_to_decoders[
-        "question"] = slim.tfexample_decoder.Tensor(
+        "question"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/question")
     data_items_to_decoders[
-        "targets"] = slim.tfexample_decoder.Tensor(
+        "targets"] = tf.contrib.slim.tfexample_decoder.Tensor(
             "image/answer")
 
     return data_fields, data_items_to_decoders

diff --git a/tensor2tensor/envs/env_problem.py b/tensor2tensor/envs/env_problem.py
@@ -33,7 +33,6 @@
 from tensor2tensor.envs import trajectory
 from tensor2tensor.layers import modalities
 import tensorflow as tf
-import tf_slim as slim
 
 # Names for data fields in stored tf.Examples.
 TIMESTEP_FIELD = "timestep"
@@ -477,7 +476,7 @@ def example_reading_spec(self):
     }
 
     data_items_to_decoders = {
-        field: slim.tfexample_decoder.Tensor(field)
+        field: tf.contrib.slim.tfexample_decoder.Tensor(field)
         for field in data_fields
     }
 

diff --git a/tensor2tensor/envs/rendered_env_problem.py b/tensor2tensor/envs/rendered_env_problem.py
@@ -25,7 +25,6 @@
 from tensor2tensor.envs import env_problem
 from tensor2tensor.envs import gym_env_problem
 import tensorflow as tf
-import tf_slim as slim
 
 _IMAGE_ENCODED_FIELD = "image/encoded"
 _IMAGE_FORMAT_FIELD = "image/format"
@@ -81,7 +80,7 @@ def example_reading_spec(self):
     # Add frame number spec and decoder.
     env_fields[_FRAME_NUMBER_FIELD] = tf.FixedLenFeature((1,), tf.int64)
     env_decoders[
-        _FRAME_NUMBER_FIELD] = slim.tfexample_decoder.Tensor(
+        _FRAME_NUMBER_FIELD] = tf.contrib.slim.tfexample_decoder.Tensor(
             _FRAME_NUMBER_FIELD)
 
     # Add video fields and decoders