Internal change

PiperOrigin-RevId: 303780351

official/vision/image_classification/augment.py

@@ -24,7 +24,7 @@
 from __future__ import print_function
 
 import math
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 from typing import Any, Dict, Iterable, List, Optional, Text, Tuple, Union
 
 from tensorflow.python.keras.layers.preprocessing import image_preprocessing as image_ops
@@ -75,8 +75,7 @@ def from_4d(image: tf.Tensor, ndims: int) -> tf.Tensor:
   return tf.reshape(image, new_shape)
 
 
-def _convert_translation_to_transform(
-    translations: Iterable[int]) -> tf.Tensor:
+def _convert_translation_to_transform(translations) -> tf.Tensor:
   """Converts translations to a projective transform.
 
   The translation matrix looks like this:
@@ -166,8 +165,7 @@ def _convert_angles_to_transform(
   )
 
 
-def transform(image: tf.Tensor,
-              transforms: Iterable[float]) -> tf.Tensor:
+def transform(image: tf.Tensor, transforms) -> tf.Tensor:
   """Prepares input data for `image_ops.transform`."""
   original_ndims = tf.rank(image)
   transforms = tf.convert_to_tensor(transforms, dtype=tf.float32)
@@ -181,8 +179,7 @@ def transform(image: tf.Tensor,
   return from_4d(image, original_ndims)
 
 
-def translate(image: tf.Tensor,
-              translations: Iterable[int]) -> tf.Tensor:
+def translate(image: tf.Tensor, translations) -> tf.Tensor:
   """Translates image(s) by provided vectors.
 
   Args:
@@ -577,7 +574,7 @@ def unwrap(image: tf.Tensor, replace: int) -> tf.Tensor:
   return image
 
 
-def _randomly_negate_tensor(tensor: tf.Tensor) -> tf.Tensor:
+def _randomly_negate_tensor(tensor):
   """With 50% prob turn the tensor negative."""
   should_flip = tf.cast(tf.floor(tf.random.uniform([]) + 0.5), tf.bool)
   final_tensor = tf.cond(should_flip, lambda: tensor, lambda: -tensor)

official/vision/image_classification/augment_test.py

@@ -21,7 +21,7 @@
 
 from absl.testing import parameterized
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 
 from official.vision.image_classification import augment
 
@@ -133,5 +133,4 @@ def test_all_policy_ops(self):
     self.assertEqual((224, 224, 3), image.shape)
 
 if __name__ == '__main__':
-  assert tf.version.VERSION.startswith('2.')
   tf.test.main()

official/vision/image_classification/classifier_trainer.py

@@ -27,7 +27,7 @@
 from absl import app
 from absl import flags
 from absl import logging
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 
 from official.modeling import performance
 from official.modeling.hyperparams import params_dict
@@ -423,5 +423,4 @@ def main(_):
   flags.mark_flag_as_required('model_type')
   flags.mark_flag_as_required('dataset')
 
-  assert tf.version.VERSION.startswith('2.')
   app.run(main)

official/vision/image_classification/classifier_trainer_test.py

@@ -30,7 +30,7 @@
 
 from absl import flags
 from absl.testing import parameterized
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 
 from tensorflow.python.distribute import combinations
 from tensorflow.python.distribute import strategy_combinations
@@ -313,5 +313,4 @@ def test_serialize_config(self):
     tf.io.gfile.rmtree(model_dir)
 
 if __name__ == '__main__':
-  assert tf.version.VERSION.startswith('2.')
   tf.test.main()

official/vision/image_classification/dataset_factory.py

@@ -23,7 +23,7 @@
 from typing import Any, List, Optional, Tuple, Mapping, Union
 from absl import logging
 from dataclasses import dataclass
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_datasets as tfds
 
 from official.modeling.hyperparams import base_config

official/vision/image_classification/efficientnet/common_modules.py

@@ -21,7 +21,7 @@
 import numpy as np
 
 import tensorflow.compat.v1 as tf1
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 from typing import Text, Optional
 
 from tensorflow.python.tpu import tpu_function

official/vision/image_classification/efficientnet/efficientnet_model.py

@@ -30,7 +30,7 @@
 
 from absl import logging
 from dataclasses import dataclass
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 
 from official.modeling import tf_utils
 from official.modeling.hyperparams import base_config

official/vision/image_classification/learning_rate.py

@@ -20,7 +20,7 @@
 
 from typing import Any, List, Mapping
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 
 BASE_LEARNING_RATE = 0.1
 

official/vision/image_classification/learning_rate_test.py

@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 
 from official.vision.image_classification import learning_rate
 
@@ -86,5 +86,4 @@ def test_piecewise_constant_decay_invalid_boundaries(self):
 
 
 if __name__ == '__main__':
-  assert tf.version.VERSION.startswith('2.')
   tf.test.main()

official/vision/image_classification/mnist_test.py

@@ -88,5 +88,4 @@ def test_end_to_end(self, distribution):
 
 
 if __name__ == "__main__":
-  tf.compat.v1.enable_v2_behavior()
   tf.test.main()

official/vision/image_classification/optimizer_factory.py

@@ -19,7 +19,7 @@
 from __future__ import print_function
 
 from absl import logging
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 import tensorflow_addons as tfa
 
 from typing import Any, Dict, Text

official/vision/image_classification/optimizer_factory_test.py

@@ -19,7 +19,7 @@
 # from __future__ import google_type_annotations
 from __future__ import print_function
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 
 from absl.testing import parameterized
 from official.vision.image_classification import optimizer_factory
@@ -111,5 +111,4 @@ def test_learning_rate_with_decay_and_warmup(self, lr_decay_type):
 
 
 if __name__ == '__main__':
-  assert tf.version.VERSION.startswith('2.')
   tf.test.main()

official/vision/image_classification/preprocessing.py

@@ -19,7 +19,7 @@
 # from __future__ import google_type_annotations
 from __future__ import print_function
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 from typing import List, Optional, Text, Tuple
 
 from official.vision.image_classification import augment

official/vision/image_classification/resnet/resnet_runnable.py

@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 
 from official.modeling import performance
 from official.staging.training import grad_utils

official/vision/image_classification/resnet/tfhub_export.py

@@ -24,7 +24,7 @@
 from absl import app
 from absl import flags
 
-import tensorflow.compat.v2 as tf
+import tensorflow as tf
 
 from official.vision.image_classification.resnet import imagenet_preprocessing
 from official.vision.image_classification.resnet import resnet_model