Revert "tf_upgrade_v2 on resnet and utils folders. (#6154)" (#6162)

This reverts commit d6b2b83c.

official/resnet/cifar10_download_and_extract.py

@@ -60,4 +60,4 @@ def main(_):
 
 if __name__ == '__main__':
   FLAGS, unparsed = parser.parse_known_args()
-  tf.compat.v1.app.run(argv=[sys.argv[0]] + unparsed)
+  tf.app.run(argv=[sys.argv[0]] + unparsed)

official/resnet/cifar10_main.py

@@ -52,7 +52,7 @@ DATASET_NAME = 'CIFAR-10'
 ###############################################################################
 def get_filenames(is_training, data_dir):
   """Returns a list of filenames."""
-  assert tf.io.gfile.exists(data_dir), (
+  assert tf.gfile.Exists(data_dir), (
       'Run cifar10_download_and_extract.py first to download and extract the '
       'CIFAR-10 data.')
 
@@ -68,7 +68,7 @@ def get_filenames(is_training, data_dir):
 def parse_record(raw_record, is_training, dtype):
   """Parse CIFAR-10 image and label from a raw record."""
   # Convert bytes to a vector of uint8 that is record_bytes long.
-  record_vector = tf.io.decode_raw(raw_record, tf.uint8)
+  record_vector = tf.decode_raw(raw_record, tf.uint8)
 
   # The first byte represents the label, which we convert from uint8 to int32
   # and then to one-hot.
@@ -81,7 +81,7 @@ def parse_record(raw_record, is_training, dtype):
 
   # Convert from [depth, height, width] to [height, width, depth], and cast as
   # float32.
-  image = tf.cast(tf.transpose(a=depth_major, perm=[1, 2, 0]), tf.float32)
+  image = tf.cast(tf.transpose(depth_major, [1, 2, 0]), tf.float32)
 
   image = preprocess_image(image, is_training)
   image = tf.cast(image, dtype)
@@ -97,7 +97,7 @@ def preprocess_image(image, is_training):
         image, HEIGHT + 8, WIDTH + 8)
 
     # Randomly crop a [HEIGHT, WIDTH] section of the image.
-    image = tf.image.random_crop(image, [HEIGHT, WIDTH, NUM_CHANNELS])
+    image = tf.random_crop(image, [HEIGHT, WIDTH, NUM_CHANNELS])
 
     # Randomly flip the image horizontally.
     image = tf.image.random_flip_left_right(image)
@@ -253,9 +253,8 @@ def run_cifar(flags_obj):
     Dictionary of results. Including final accuracy.
   """
   if flags_obj.image_bytes_as_serving_input:
-    tf.compat.v1.logging.fatal(
-        '--image_bytes_as_serving_input cannot be set to True for CIFAR. '
-        'This flag is only applicable to ImageNet.')
+    tf.logging.fatal('--image_bytes_as_serving_input cannot be set to True '
+                     'for CIFAR. This flag is only applicable to ImageNet.')
     return
 
   input_function = (flags_obj.use_synthetic_data and
@@ -274,6 +273,6 @@ def main(_):
 
 
 if __name__ == '__main__':
-  tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)
+  tf.logging.set_verbosity(tf.logging.INFO)
   define_cifar_flags()
   absl_app.run(main)

official/resnet/cifar10_test.py

@@ -25,7 +25,7 @@ import tensorflow as tf  # pylint: disable=g-bad-import-order
 from official.resnet import cifar10_main
 from official.utils.testing import integration
 
-tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+tf.logging.set_verbosity(tf.logging.ERROR)
 
 _BATCH_SIZE = 128
 _HEIGHT = 32
@@ -44,7 +44,7 @@ class BaseTest(tf.test.TestCase):
 
   def tearDown(self):
     super(BaseTest, self).tearDown()
-    tf.io.gfile.rmtree(self.get_temp_dir())
+    tf.gfile.DeleteRecursively(self.get_temp_dir())
 
   def test_dataset_input_fn(self):
     fake_data = bytearray()
@@ -62,8 +62,7 @@ class BaseTest(tf.test.TestCase):
         filename, cifar10_main._RECORD_BYTES)  # pylint: disable=protected-access
     fake_dataset = fake_dataset.map(
         lambda val: cifar10_main.parse_record(val, False, tf.float32))
-    image, label = tf.compat.v1.data.make_one_shot_iterator(
-        fake_dataset).get_next()
+    image, label = fake_dataset.make_one_shot_iterator().get_next()
 
     self.assertAllEqual(label.shape, ())
     self.assertAllEqual(image.shape, (_HEIGHT, _WIDTH, _NUM_CHANNELS))
@@ -80,7 +79,7 @@ class BaseTest(tf.test.TestCase):
   def cifar10_model_fn_helper(self, mode, resnet_version, dtype):
     input_fn = cifar10_main.get_synth_input_fn(dtype)
     dataset = input_fn(True, '', _BATCH_SIZE)
-    iterator = tf.compat.v1.data.make_initializable_iterator(dataset)
+    iterator = dataset.make_initializable_iterator()
     features, labels = iterator.get_next()
     spec = cifar10_main.cifar10_model_fn(
         features, labels, mode, {
@@ -143,7 +142,7 @@ class BaseTest(tf.test.TestCase):
     model = cifar10_main.Cifar10Model(32, data_format='channels_last',
                                       num_classes=num_classes,
                                       resnet_version=resnet_version)
-    fake_input = tf.random.uniform([batch_size, _HEIGHT, _WIDTH, _NUM_CHANNELS])
+    fake_input = tf.random_uniform([batch_size, _HEIGHT, _WIDTH, _NUM_CHANNELS])
     output = model(fake_input, training=True)
 
     self.assertAllEqual(output.shape, (batch_size, num_classes))

official/resnet/estimator_cifar_benchmark.py

@@ -144,7 +144,7 @@ class EstimatorCifar10BenchmarkTests(tf.test.Benchmark):
     return os.path.join(self.output_dir, folder_name)
 
   def _setup(self):
-    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.DEBUG)
+    tf.logging.set_verbosity(tf.logging.DEBUG)
     if EstimatorCifar10BenchmarkTests.local_flags is None:
       cifar_main.define_cifar_flags()
       # Loads flags to get defaults to then override.

official/resnet/imagenet_main.py

@@ -95,14 +95,14 @@ def _parse_example_proto(example_serialized):
   """
   # Dense features in Example proto.
   feature_map = {
-      'image/encoded': tf.io.FixedLenFeature([], dtype=tf.string,
+      'image/encoded': tf.FixedLenFeature([], dtype=tf.string,
                                           default_value=''),
-      'image/class/label': tf.io.FixedLenFeature([], dtype=tf.int64,
+      'image/class/label': tf.FixedLenFeature([], dtype=tf.int64,
                                               default_value=-1),
-      'image/class/text': tf.io.FixedLenFeature([], dtype=tf.string,
+      'image/class/text': tf.FixedLenFeature([], dtype=tf.string,
                                              default_value=''),
   }
-  sparse_float32 = tf.io.VarLenFeature(dtype=tf.float32)
+  sparse_float32 = tf.VarLenFeature(dtype=tf.float32)
   # Sparse features in Example proto.
   feature_map.update(
       {k: sparse_float32 for k in ['image/object/bbox/xmin',
@@ -110,8 +110,7 @@ def _parse_example_proto(example_serialized):
                                    'image/object/bbox/xmax',
                                    'image/object/bbox/ymax']})
 
-  features = tf.io.parse_single_example(serialized=example_serialized,
-                                        features=feature_map)
+  features = tf.parse_single_example(example_serialized, feature_map)
   label = tf.cast(features['image/class/label'], dtype=tf.int32)
 
   xmin = tf.expand_dims(features['image/object/bbox/xmin'].values, 0)
@@ -125,7 +124,7 @@ def _parse_example_proto(example_serialized):
   # Force the variable number of bounding boxes into the shape
   # [1, num_boxes, coords].
   bbox = tf.expand_dims(bbox, 0)
-  bbox = tf.transpose(a=bbox, perm=[0, 2, 1])
+  bbox = tf.transpose(bbox, [0, 2, 1])
 
   return features['image/encoded'], label, bbox
 
@@ -189,7 +188,7 @@ def input_fn(is_training, data_dir, batch_size, num_epochs=1,
   # This number is low enough to not cause too much contention on small systems
   # but high enough to provide the benefits of parallelization. You may want
   # to increase this number if you have a large number of CPU cores.
-  dataset = dataset.apply(tf.data.experimental.parallel_interleave(
+  dataset = dataset.apply(tf.contrib.data.parallel_interleave(
       tf.data.TFRecordDataset, cycle_length=10))
 
   return resnet_run_loop.process_record_dataset(
@@ -353,6 +352,6 @@ def main(_):
 
 
 if __name__ == '__main__':
-  tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)
+  tf.logging.set_verbosity(tf.logging.INFO)
   define_imagenet_flags()
   absl_app.run(main)

official/resnet/imagenet_preprocessing.py

@@ -108,7 +108,7 @@ def _central_crop(image, crop_height, crop_width):
   Returns:
     3-D tensor with cropped image.
   """
-  shape = tf.shape(input=image)
+  shape = tf.shape(image)
   height, width = shape[0], shape[1]
 
   amount_to_be_cropped_h = (height - crop_height)
@@ -195,7 +195,7 @@ def _aspect_preserving_resize(image, resize_min):
   Returns:
     resized_image: A 3-D tensor containing the resized image.
   """
-  shape = tf.shape(input=image)
+  shape = tf.shape(image)
   height, width = shape[0], shape[1]
 
   new_height, new_width = _smallest_size_at_least(height, width, resize_min)
@@ -218,7 +218,7 @@ def _resize_image(image, height, width):
     resized_image: A 3-D tensor containing the resized image. The first two
       dimensions have the shape [height, width].
   """
-  return tf.image.resize(
+  return tf.image.resize_images(
       image, [height, width], method=tf.image.ResizeMethod.BILINEAR,
       align_corners=False)
 

official/resnet/imagenet_test.py

@@ -24,7 +24,7 @@ import tensorflow as tf  # pylint: disable=g-bad-import-order
 from official.resnet import imagenet_main
 from official.utils.testing import integration
 
-tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+tf.logging.set_verbosity(tf.logging.ERROR)
 
 _BATCH_SIZE = 32
 _LABEL_CLASSES = 1001
@@ -39,7 +39,7 @@ class BaseTest(tf.test.TestCase):
 
   def tearDown(self):
     super(BaseTest, self).tearDown()
-    tf.io.gfile.rmtree(self.get_temp_dir())
+    tf.gfile.DeleteRecursively(self.get_temp_dir())
 
   def _tensor_shapes_helper(self, resnet_size, resnet_version, dtype, with_gpu):
     """Checks the tensor shapes after each phase of the ResNet model."""
@@ -62,7 +62,7 @@ class BaseTest(tf.test.TestCase):
           resnet_version=resnet_version,
           dtype=dtype
       )
-      inputs = tf.random.uniform([1, 224, 224, 3])
+      inputs = tf.random_uniform([1, 224, 224, 3])
       output = model(inputs, training=True)
 
       initial_conv = graph.get_tensor_by_name('resnet_model/initial_conv:0')
@@ -189,11 +189,11 @@ class BaseTest(tf.test.TestCase):
 
   def resnet_model_fn_helper(self, mode, resnet_version, dtype):
     """Tests that the EstimatorSpec is given the appropriate arguments."""
-    tf.compat.v1.train.create_global_step()
+    tf.train.create_global_step()
 
     input_fn = imagenet_main.get_synth_input_fn(dtype)
     dataset = input_fn(True, '', _BATCH_SIZE)
-    iterator = tf.compat.v1.data.make_initializable_iterator(dataset)
+    iterator = dataset.make_initializable_iterator()
     features, labels = iterator.get_next()
     spec = imagenet_main.imagenet_model_fn(
         features, labels, mode, {
@@ -257,7 +257,7 @@ class BaseTest(tf.test.TestCase):
         50, data_format='channels_last', num_classes=num_classes,
         resnet_version=resnet_version)
 
-    fake_input = tf.random.uniform([batch_size, 224, 224, 3])
+    fake_input = tf.random_uniform([batch_size, 224, 224, 3])
     output = model(fake_input, training=True)
 
     self.assertAllEqual(output.shape, (batch_size, num_classes))

official/resnet/keras/keras_benchmark.py

@@ -43,7 +43,7 @@ class KerasBenchmark(tf.test.Benchmark):
 
   def _setup(self):
     """Sets up and resets flags before each test."""
-    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.DEBUG)
+    tf.logging.set_verbosity(tf.logging.DEBUG)
     if KerasBenchmark.local_flags is None:
       for flag_method in self.flag_methods:
         flag_method()

official/resnet/keras/keras_cifar_main.py

@@ -81,7 +81,7 @@ def parse_record_keras(raw_record, is_training, dtype):
     Tuple with processed image tensor and one-hot-encoded label tensor.
   """
   image, label = cifar_main.parse_record(raw_record, is_training, dtype)
-  label = tf.compat.v1.sparse_to_dense(label, (cifar_main.NUM_CLASSES,), 1)
+  label = tf.sparse_to_dense(label, (cifar_main.NUM_CLASSES,), 1)
   return image, label
 
 
@@ -98,7 +98,7 @@ def run(flags_obj):
     Dictionary of training and eval stats.
   """
   if flags_obj.enable_eager:
-    tf.compat.v1.enable_eager_execution()
+    tf.enable_eager_execution()
 
   dtype = flags_core.get_tf_dtype(flags_obj)
   if dtype == 'fp16':
@@ -194,7 +194,7 @@ def main(_):
 
 
 if __name__ == '__main__':
-  tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)
+  tf.logging.set_verbosity(tf.logging.INFO)
   cifar_main.define_cifar_flags()
   keras_common.define_keras_flags()
   absl_app.run(main)

official/resnet/keras/keras_common.py

@@ -80,8 +80,7 @@ class TimeHistory(tf.keras.callbacks.Callback):
       if batch != 0:
         self.record_batch = True
         self.timestamp_log.append(BatchTimestamp(batch, timestamp))
-        tf.compat.v1.logging.info(
-            "BenchmarkMetric: {'num_batches':%d, 'time_taken': %f,"
+        tf.logging.info("BenchmarkMetric: {'num_batches':%d, 'time_taken': %f,"
                         "'images_per_second': %f}" %
                         (batch, elapsed_time, examples_per_second))
 
@@ -121,8 +120,7 @@ class LearningRateBatchScheduler(tf.keras.callbacks.Callback):
     if lr != self.prev_lr:
       self.model.optimizer.learning_rate = lr  # lr should be a float here
       self.prev_lr = lr
-      tf.compat.v1.logging.debug(
-          'Epoch %05d Batch %05d: LearningRateBatchScheduler '
+      tf.logging.debug('Epoch %05d Batch %05d: LearningRateBatchScheduler '
                        'change learning rate to %s.', self.epochs, batch, lr)
 
 
@@ -228,20 +226,22 @@ def get_synth_input_fn(height, width, num_channels, num_classes,
   def input_fn(is_training, data_dir, batch_size, *args, **kwargs):
     """Returns dataset filled with random data."""
     # Synthetic input should be within [0, 255].
-    inputs = tf.random.truncated_normal([height, width, num_channels],
+    inputs = tf.truncated_normal(
+        [height, width, num_channels],
         dtype=dtype,
         mean=127,
         stddev=60,
         name='synthetic_inputs')
 
-    labels = tf.random.uniform([1],
+    labels = tf.random_uniform(
+        [1],
         minval=0,
         maxval=num_classes - 1,
         dtype=tf.int32,
         name='synthetic_labels')
     data = tf.data.Dataset.from_tensors((inputs, labels)).repeat()
     data = data.batch(batch_size)
-    data = data.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+    data = data.prefetch(buffer_size=tf.contrib.data.AUTOTUNE)
     return data
 
   return input_fn

official/resnet/keras/keras_common_test.py

@@ -22,7 +22,7 @@ import tensorflow as tf  # pylint: disable=g-bad-import-order
 
 from official.resnet.keras import keras_common
 
-tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+tf.logging.set_verbosity(tf.logging.ERROR)
 
 
 class KerasCommonTests(tf.test.TestCase):

official/resnet/keras/keras_imagenet_main.py

@@ -88,7 +88,7 @@ def run(flags_obj):
     ValueError: If fp16 is passed as it is not currently supported.
   """
   if flags_obj.enable_eager:
-    tf.compat.v1.enable_eager_execution()
+    tf.enable_eager_execution()
 
   dtype = flags_core.get_tf_dtype(flags_obj)
   if dtype == 'fp16':
@@ -187,7 +187,7 @@ def main(_):
 
 
 if __name__ == '__main__':
-  tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)
+  tf.logging.set_verbosity(tf.logging.INFO)
   imagenet_main.define_imagenet_flags()
   keras_common.define_keras_flags()
   absl_app.run(main)

official/resnet/layer_test.py

@@ -72,10 +72,10 @@ class BaseTest(reference_data.BaseTest):
 
     g = tf.Graph()
     with g.as_default():
-      tf.compat.v1.set_random_seed(self.name_to_seed(name))
-      input_tensor = tf.compat.v1.get_variable(
+      tf.set_random_seed(self.name_to_seed(name))
+      input_tensor = tf.get_variable(
           "input_tensor", dtype=tf.float32,
-          initializer=tf.random.uniform((32, 16, 16, 3), maxval=1)
+          initializer=tf.random_uniform((32, 16, 16, 3), maxval=1)
       )
       layer = resnet_model.batch_norm(
           inputs=input_tensor, data_format=DATA_FORMAT, training=True)
@@ -137,7 +137,7 @@ class BaseTest(reference_data.BaseTest):
 
     g = tf.Graph()
     with g.as_default():
-      tf.compat.v1.set_random_seed(self.name_to_seed(name))
+      tf.set_random_seed(self.name_to_seed(name))
       strides = 1
       channels_out = channels
       projection_shortcut = None
@@ -151,9 +151,9 @@ class BaseTest(reference_data.BaseTest):
       if bottleneck:
         filters = channels_out // 4
 
-      input_tensor = tf.compat.v1.get_variable(
+      input_tensor = tf.get_variable(
           "input_tensor", dtype=tf.float32,
-          initializer=tf.random.uniform((batch_size, width, width, channels),
+          initializer=tf.random_uniform((batch_size, width, width, channels),
                                         maxval=1)
       )
 

official/resnet/resnet_model.py

@@ -48,7 +48,7 @@ def batch_norm(inputs, training, data_format):
   """Performs a batch normalization using a standard set of parameters."""
   # We set fused=True for a significant performance boost. See
   # https://www.tensorflow.org/performance/performance_guide#common_fused_ops
-  return tf.compat.v1.layers.batch_normalization(
+  return tf.layers.batch_normalization(
       inputs=inputs, axis=1 if data_format == 'channels_first' else 3,
       momentum=_BATCH_NORM_DECAY, epsilon=_BATCH_NORM_EPSILON, center=True,
       scale=True, training=training, fused=True)
@@ -73,12 +73,10 @@ def fixed_padding(inputs, kernel_size, data_format):
   pad_end = pad_total - pad_beg
 
   if data_format == 'channels_first':
-    padded_inputs = tf.pad(tensor=inputs,
-                           paddings=[[0, 0], [0, 0], [pad_beg, pad_end],
-                                     [pad_beg, pad_end]])
+    padded_inputs = tf.pad(inputs, [[0, 0], [0, 0],
+                                    [pad_beg, pad_end], [pad_beg, pad_end]])
   else:
-    padded_inputs = tf.pad(tensor=inputs,
-                           paddings=[[0, 0], [pad_beg, pad_end],
+    padded_inputs = tf.pad(inputs, [[0, 0], [pad_beg, pad_end],
                                     [pad_beg, pad_end], [0, 0]])
   return padded_inputs
 
@@ -90,10 +88,10 @@ def conv2d_fixed_padding(inputs, filters, kernel_size, strides, data_format):
   if strides > 1:
     inputs = fixed_padding(inputs, kernel_size, data_format)
 
-  return tf.compat.v1.layers.conv2d(
+  return tf.layers.conv2d(
       inputs=inputs, filters=filters, kernel_size=kernel_size, strides=strides,
       padding=('SAME' if strides == 1 else 'VALID'), use_bias=False,
-      kernel_initializer=tf.compat.v1.variance_scaling_initializer(),
+      kernel_initializer=tf.variance_scaling_initializer(),
       data_format=data_format)
 
 
@@ -477,7 +475,7 @@ class Model(object):
       A variable scope for the model.
     """
 
-    return tf.compat.v1.variable_scope('resnet_model',
+    return tf.variable_scope('resnet_model',
                              custom_getter=self._custom_dtype_getter)
 
   def __call__(self, inputs, training):
@@ -497,7 +495,7 @@ class Model(object):
         # Convert the inputs from channels_last (NHWC) to channels_first (NCHW).
         # This provides a large performance boost on GPU. See
         # https://www.tensorflow.org/performance/performance_guide#data_formats
-        inputs = tf.transpose(a=inputs, perm=[0, 3, 1, 2])
+        inputs = tf.transpose(inputs, [0, 3, 1, 2])
 
       inputs = conv2d_fixed_padding(
           inputs=inputs, filters=self.num_filters, kernel_size=self.kernel_size,
@@ -513,7 +511,7 @@ class Model(object):
         inputs = tf.nn.relu(inputs)
 
       if self.first_pool_size:
-        inputs = tf.compat.v1.layers.max_pooling2d(
+        inputs = tf.layers.max_pooling2d(
             inputs=inputs, pool_size=self.first_pool_size,
             strides=self.first_pool_stride, padding='SAME',
             data_format=self.data_format)
@@ -539,10 +537,10 @@ class Model(object):
       # but that is the same as doing a reduce_mean. We do a reduce_mean
       # here because it performs better than AveragePooling2D.
       axes = [2, 3] if self.data_format == 'channels_first' else [1, 2]
-      inputs = tf.reduce_mean(input_tensor=inputs, axis=axes, keepdims=True)
+      inputs = tf.reduce_mean(inputs, axes, keepdims=True)
       inputs = tf.identity(inputs, 'final_reduce_mean')
 
       inputs = tf.squeeze(inputs, axes)
-      inputs = tf.compat.v1.layers.dense(inputs=inputs, units=self.num_classes)
+      inputs = tf.layers.dense(inputs=inputs, units=self.num_classes)
       inputs = tf.identity(inputs, 'final_dense')
       return inputs

official/resnet/resnet_run_loop.py

@@ -88,7 +88,7 @@ def process_record_dataset(dataset,
 
   # Parses the raw records into images and labels.
   dataset = dataset.apply(
-      tf.data.experimental.map_and_batch(
+      tf.contrib.data.map_and_batch(
           lambda value: parse_record_fn(value, is_training, dtype),
           batch_size=batch_size,
           num_parallel_batches=num_parallel_batches,
@@ -100,11 +100,11 @@ def process_record_dataset(dataset,
   # critical training path. Setting buffer_size to tf.contrib.data.AUTOTUNE
   # allows DistributionStrategies to adjust how many batches to fetch based
   # on how many devices are present.
-  dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+  dataset = dataset.prefetch(buffer_size=tf.contrib.data.AUTOTUNE)
 
   # Defines a specific size thread pool for tf.data operations.
   if datasets_num_private_threads:
-    tf.compat.v1.logging.info('datasets_num_private_threads: %s',
+    tf.logging.info('datasets_num_private_threads: %s',
                     datasets_num_private_threads)
     dataset = threadpool.override_threadpool(
         dataset,
@@ -140,21 +140,21 @@ def get_synth_input_fn(height, width, num_channels, num_classes,
   def input_fn(is_training, data_dir, batch_size, *args, **kwargs):
     """Returns dataset filled with random data."""
     # Synthetic input should be within [0, 255].
-    inputs = tf.random.truncated_normal(
+    inputs = tf.truncated_normal(
         [batch_size] + [height, width, num_channels],
         dtype=dtype,
         mean=127,
         stddev=60,
         name='synthetic_inputs')
 
-    labels = tf.random.uniform(
+    labels = tf.random_uniform(
         [batch_size],
         minval=0,
         maxval=num_classes - 1,
         dtype=tf.int32,
         name='synthetic_labels')
     data = tf.data.Dataset.from_tensors((inputs, labels)).repeat()
-    data = data.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+    data = data.prefetch(buffer_size=tf.contrib.data.AUTOTUNE)
     return data
 
   return input_fn
@@ -172,7 +172,7 @@ def image_bytes_serving_input_fn(image_shape, dtype=tf.float32):
         image_bytes, bbox, height, width, num_channels, is_training=False)
     return image
 
-  image_bytes_list = tf.compat.v1.placeholder(
+  image_bytes_list = tf.placeholder(
       shape=[None], dtype=tf.string, name='input_tensor')
   images = tf.map_fn(
       _preprocess_image, image_bytes_list, back_prop=False, dtype=dtype)
@@ -197,17 +197,15 @@ def override_flags_and_set_envars_for_gpu_thread_pool(flags_obj):
     what has been set by the user on the command-line.
   """
   cpu_count = multiprocessing.cpu_count()
-  tf.compat.v1.logging.info('Logical CPU cores: %s', cpu_count)
+  tf.logging.info('Logical CPU cores: %s', cpu_count)
 
   # Sets up thread pool for each GPU for op scheduling.
   per_gpu_thread_count = 1
   total_gpu_thread_count = per_gpu_thread_count * flags_obj.num_gpus
   os.environ['TF_GPU_THREAD_MODE'] = flags_obj.tf_gpu_thread_mode
   os.environ['TF_GPU_THREAD_COUNT'] = str(per_gpu_thread_count)
-  tf.compat.v1.logging.info('TF_GPU_THREAD_COUNT: %s',
-                            os.environ['TF_GPU_THREAD_COUNT'])
-  tf.compat.v1.logging.info('TF_GPU_THREAD_MODE: %s',
-                            os.environ['TF_GPU_THREAD_MODE'])
+  tf.logging.info('TF_GPU_THREAD_COUNT: %s', os.environ['TF_GPU_THREAD_COUNT'])
+  tf.logging.info('TF_GPU_THREAD_MODE: %s', os.environ['TF_GPU_THREAD_MODE'])
 
   # Reduces general thread pool by number of threads used for GPU pool.
   main_thread_count = cpu_count - total_gpu_thread_count
@@ -258,15 +256,13 @@ def learning_rate_with_decay(
 
   def learning_rate_fn(global_step):
     """Builds scaled learning rate function with 5 epoch warm up."""
-    lr = tf.compat.v1.train.piecewise_constant(global_step, boundaries, vals)
+    lr = tf.train.piecewise_constant(global_step, boundaries, vals)
     if warmup:
       warmup_steps = int(batches_per_epoch * 5)
       warmup_lr = (
           initial_learning_rate * tf.cast(global_step, tf.float32) / tf.cast(
               warmup_steps, tf.float32))
-      return tf.cond(pred=global_step < warmup_steps,
-                     true_fn=lambda: warmup_lr,
-                     false_fn=lambda: lr)
+      return tf.cond(global_step < warmup_steps, lambda: warmup_lr, lambda: lr)
     return lr
 
   return learning_rate_fn
@@ -317,7 +313,7 @@ def resnet_model_fn(features, labels, mode, model_class,
   """
 
   # Generate a summary node for the images
-  tf.compat.v1.summary.image('images', features, max_outputs=6)
+  tf.summary.image('images', features, max_outputs=6)
   # Checks that features/images have same data type being used for calculations.
   assert features.dtype == dtype
 
@@ -332,7 +328,7 @@ def resnet_model_fn(features, labels, mode, model_class,
   logits = tf.cast(logits, tf.float32)
 
   predictions = {
-      'classes': tf.argmax(input=logits, axis=1),
+      'classes': tf.argmax(logits, axis=1),
       'probabilities': tf.nn.softmax(logits, name='softmax_tensor')
   }
 
@@ -346,12 +342,12 @@ def resnet_model_fn(features, labels, mode, model_class,
         })
 
   # Calculate loss, which includes softmax cross entropy and L2 regularization.
-  cross_entropy = tf.compat.v1.losses.sparse_softmax_cross_entropy(
+  cross_entropy = tf.losses.sparse_softmax_cross_entropy(
       logits=logits, labels=labels)
 
   # Create a tensor named cross_entropy for logging purposes.
   tf.identity(cross_entropy, name='cross_entropy')
-  tf.compat.v1.summary.scalar('cross_entropy', cross_entropy)
+  tf.summary.scalar('cross_entropy', cross_entropy)
 
   # If no loss_filter_fn is passed, assume we want the default behavior,
   # which is that batch_normalization variables are excluded from loss.
@@ -362,21 +358,21 @@ def resnet_model_fn(features, labels, mode, model_class,
   # Add weight decay to the loss.
   l2_loss = weight_decay * tf.add_n(
       # loss is computed using fp32 for numerical stability.
-      [tf.nn.l2_loss(tf.cast(v, tf.float32))
-       for v in tf.compat.v1.trainable_variables() if loss_filter_fn(v.name)])
-  tf.compat.v1.summary.scalar('l2_loss', l2_loss)
+      [tf.nn.l2_loss(tf.cast(v, tf.float32)) for v in tf.trainable_variables()
+       if loss_filter_fn(v.name)])
+  tf.summary.scalar('l2_loss', l2_loss)
   loss = cross_entropy + l2_loss
 
   if mode == tf.estimator.ModeKeys.TRAIN:
-    global_step = tf.compat.v1.train.get_or_create_global_step()
+    global_step = tf.train.get_or_create_global_step()
 
     learning_rate = learning_rate_fn(global_step)
 
     # Create a tensor named learning_rate for logging purposes
     tf.identity(learning_rate, name='learning_rate')
-    tf.compat.v1.summary.scalar('learning_rate', learning_rate)
+    tf.summary.scalar('learning_rate', learning_rate)
 
-    optimizer = tf.compat.v1.train.MomentumOptimizer(
+    optimizer = tf.train.MomentumOptimizer(
         learning_rate=learning_rate,
         momentum=momentum
     )
@@ -413,22 +409,24 @@ def resnet_model_fn(features, labels, mode, model_class,
         grad_vars = _dense_grad_filter(grad_vars)
       minimize_op = optimizer.apply_gradients(grad_vars, global_step)
 
-    update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
+    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
     train_op = tf.group(minimize_op, update_ops)
   else:
     train_op = None
 
-  accuracy = tf.compat.v1.metrics.accuracy(labels, predictions['classes'])
-  accuracy_top_5 = tf.compat.v1.metrics.mean(
-      tf.nn.in_top_k(predictions=logits, targets=labels, k=5, name='top_5_op'))
+  accuracy = tf.metrics.accuracy(labels, predictions['classes'])
+  accuracy_top_5 = tf.metrics.mean(tf.nn.in_top_k(predictions=logits,
+                                                  targets=labels,
+                                                  k=5,
+                                                  name='top_5_op'))
   metrics = {'accuracy': accuracy,
              'accuracy_top_5': accuracy_top_5}
 
   # Create a tensor named train_accuracy for logging purposes
   tf.identity(accuracy[1], name='train_accuracy')
   tf.identity(accuracy_top_5[1], name='train_accuracy_top_5')
-  tf.compat.v1.summary.scalar('train_accuracy', accuracy[1])
-  tf.compat.v1.summary.scalar('train_accuracy_top_5', accuracy_top_5[1])
+  tf.summary.scalar('train_accuracy', accuracy[1])
+  tf.summary.scalar('train_accuracy_top_5', accuracy_top_5[1])
 
   return tf.estimator.EstimatorSpec(
       mode=mode,
@@ -467,7 +465,7 @@ def resnet_main(
 
   # Creates session config. allow_soft_placement = True, is required for
   # multi-GPU and is not harmful for other modes.
-  session_config = tf.compat.v1.ConfigProto(
+  session_config = tf.ConfigProto(
       inter_op_parallelism_threads=flags_obj.inter_op_parallelism_threads,
       intra_op_parallelism_threads=flags_obj.intra_op_parallelism_threads,
       allow_soft_placement=True)
@@ -559,14 +557,13 @@ def resnet_main(
     schedule[-1] = flags_obj.train_epochs - sum(schedule[:-1])  # over counting.
 
   for cycle_index, num_train_epochs in enumerate(schedule):
-    tf.compat.v1.logging.info('Starting cycle: %d/%d', cycle_index,
-                              int(n_loops))
+    tf.logging.info('Starting cycle: %d/%d', cycle_index, int(n_loops))
 
     if num_train_epochs:
       classifier.train(input_fn=lambda: input_fn_train(num_train_epochs),
                        hooks=train_hooks, max_steps=flags_obj.max_train_steps)
 
-    tf.compat.v1.logging.info('Starting to evaluate.')
+    tf.logging.info('Starting to evaluate.')
 
     # flags_obj.max_train_steps is generally associated with testing and
     # profiling. As a result it is frequently called with synthetic data, which

official/utils/accelerator/tpu.py

@@ -71,7 +71,7 @@ def construct_scalar_host_call(metric_dict, model_dir, prefix=""):
   # expects [batch_size, ...] Tensors, thus reshape to introduce a batch
   # dimension. These Tensors are implicitly concatenated to
   # [params['batch_size']].
-  global_step_tensor = tf.reshape(tf.compat.v1.train.get_or_create_global_step(), [1])
+  global_step_tensor = tf.reshape(tf.train.get_or_create_global_step(), [1])
   other_tensors = [tf.reshape(metric_dict[key], [1]) for key in metric_names]
 
   return host_call_fn, [global_step_tensor] + other_tensors

official/utils/accelerator/tpu_test.py

@@ -36,14 +36,14 @@ class TPUBaseTester(tf.test.TestCase):
     np.random.seed(seed)
 
     embeddings = np.random.random(size=(vocab_size, embedding_dim))
-    embedding_table = tf.convert_to_tensor(value=embeddings, dtype=tf.float32)
+    embedding_table = tf.convert_to_tensor(embeddings, dtype=tf.float32)
 
     tokens = np.random.randint(low=1, high=vocab_size-1,
                                size=(batch_size, sequence_length))
     for i in range(batch_size):
       tokens[i, np.random.randint(low=0, high=sequence_length-1):] = 0
-    values = tf.convert_to_tensor(value=tokens, dtype=tf.int32)
-    mask = tf.cast(tf.not_equal(values, 0), dtype=tf.float32)
+    values = tf.convert_to_tensor(tokens, dtype=tf.int32)
+    mask = tf.to_float(tf.not_equal(values, 0))
     return embedding_table, values, mask
 
   def _test_embedding(self, embedding_dim, vocab_size,

official/utils/data/file_io.py

@@ -47,11 +47,11 @@ class _GarbageCollector(object):
   def purge(self):
     try:
       for i in self.temp_buffers:
-        if tf.io.gfile.exists(i):
-          tf.io.gfile.remove(i)
-          tf.compat.v1.logging.info("Buffer file {} removed".format(i))
+        if tf.gfile.Exists(i):
+          tf.gfile.Remove(i)
+          tf.logging.info("Buffer file {} removed".format(i))
     except Exception as e:
-      tf.compat.v1.logging.error("Failed to cleanup buffer files: {}".format(e))
+      tf.logging.error("Failed to cleanup buffer files: {}".format(e))
 
 
 _GARBAGE_COLLECTOR = _GarbageCollector()
@@ -64,7 +64,7 @@ def write_to_temp_buffer(dataframe, buffer_folder, columns):
   if buffer_folder is None:
     _, buffer_path = tempfile.mkstemp()
   else:
-    tf.io.gfile.makedirs(buffer_folder)
+    tf.gfile.MakeDirs(buffer_folder)
     buffer_path = os.path.join(buffer_folder, str(uuid.uuid4()))
   _GARBAGE_COLLECTOR.register(buffer_path)
 
@@ -169,35 +169,35 @@ def write_to_buffer(dataframe, buffer_path, columns, expected_size=None):
   Returns:
     The path of the buffer.
   """
-  if tf.io.gfile.exists(buffer_path) and tf.io.gfile.stat(buffer_path).length > 0:
-    actual_size = tf.io.gfile.stat(buffer_path).length
+  if tf.gfile.Exists(buffer_path) and tf.gfile.Stat(buffer_path).length > 0:
+    actual_size = tf.gfile.Stat(buffer_path).length
     if expected_size == actual_size:
       return buffer_path
-    tf.compat.v1.logging.warning(
+    tf.logging.warning(
         "Existing buffer {} has size {}. Expected size {}. Deleting and "
         "rebuilding buffer.".format(buffer_path, actual_size, expected_size))
-    tf.io.gfile.remove(buffer_path)
+    tf.gfile.Remove(buffer_path)
 
   if dataframe is None:
     raise ValueError(
         "dataframe was None but a valid existing buffer was not found.")
 
-  tf.io.gfile.makedirs(os.path.split(buffer_path)[0])
+  tf.gfile.MakeDirs(os.path.split(buffer_path)[0])
 
-  tf.compat.v1.logging.info("Constructing TFRecordDataset buffer: {}".format(buffer_path))
+  tf.logging.info("Constructing TFRecordDataset buffer: {}".format(buffer_path))
 
   count = 0
   pool = multiprocessing.Pool(multiprocessing.cpu_count())
   try:
-    with tf.io.TFRecordWriter(buffer_path) as writer:
+    with tf.python_io.TFRecordWriter(buffer_path) as writer:
       for df_shards in iter_shard_dataframe(df=dataframe,
                                             rows_per_core=_ROWS_PER_CORE):
         _serialize_shards(df_shards, columns, pool, writer)
         count += sum([len(s) for s in df_shards])
-        tf.compat.v1.logging.info("{}/{} examples written."
+        tf.logging.info("{}/{} examples written."
                         .format(str(count).ljust(8), len(dataframe)))
   finally:
     pool.terminate()
 
-  tf.compat.v1.logging.info("Buffer write complete.")
+  tf.logging.info("Buffer write complete.")
   return buffer_path

official/utils/data/file_io_test.py

@@ -77,9 +77,9 @@ _TEST_CASES = [
 ]
 
 _FEATURE_MAP = {
-    _RAW_ROW: tf.io.FixedLenFeature([1], dtype=tf.int64),
-    _DUMMY_COL: tf.io.FixedLenFeature([1], dtype=tf.int64),
-    _DUMMY_VEC_COL: tf.io.FixedLenFeature([_DUMMY_VEC_LEN], dtype=tf.float32)
+    _RAW_ROW: tf.FixedLenFeature([1], dtype=tf.int64),
+    _DUMMY_COL: tf.FixedLenFeature([1], dtype=tf.int64),
+    _DUMMY_VEC_COL: tf.FixedLenFeature([_DUMMY_VEC_LEN], dtype=tf.float32)
 }
 
 
@@ -156,9 +156,9 @@ class BaseTest(tf.test.TestCase):
     with self.test_session(graph=tf.Graph()) as sess:
       dataset = tf.data.TFRecordDataset(buffer_path)
       dataset = dataset.batch(1).map(
-          lambda x: tf.io.parse_example(serialized=x, features=_FEATURE_MAP))
+          lambda x: tf.parse_example(x, _FEATURE_MAP))
 
-      data_iter = tf.compat.v1.data.make_one_shot_iterator(dataset)
+      data_iter = dataset.make_one_shot_iterator()
       seen_rows = set()
       for i in range(num_rows+5):
         row = data_iter.get_next()
@@ -177,7 +177,7 @@ class BaseTest(tf.test.TestCase):
           self.assertGreaterEqual(i, num_rows, msg="Too few rows.")
 
     file_io._GARBAGE_COLLECTOR.purge()
-    assert not tf.io.gfile.exists(buffer_path)
+    assert not tf.gfile.Exists(buffer_path)
 
   def test_serialize_deserialize_0(self):
     self._serialize_deserialize(num_cores=1)

official/utils/export/export.py

@@ -40,7 +40,7 @@ def build_tensor_serving_input_receiver_fn(shape, dtype=tf.float32,
   """
   def serving_input_receiver_fn():
     # Prep a placeholder where the input example will be fed in
-    features = tf.compat.v1.placeholder(
+    features = tf.placeholder(
         dtype=dtype, shape=[batch_size] + shape, name='input_tensor')
 
     return tf.estimator.export.TensorServingInputReceiver(