Changes needed for T2T's problems to be imported on TF 2.0 -- this doesn't mean we are using these problems.

This is only so that we can `import trax` on these environments.

PiperOrigin-RevId: 282108974

Author: afrozenator

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={

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),

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
 

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

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)

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):

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

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)

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",

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)