diff --git a/official/vision/image_classification/augment.py b/official/vision/image_classification/augment.py index a71c8d00832..4ba8bb6e4b8 100644 --- a/official/vision/image_classification/augment.py +++ b/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) diff --git a/official/vision/image_classification/augment_test.py b/official/vision/image_classification/augment_test.py index 40c152e4f3d..58fbac79246 100644 --- a/official/vision/image_classification/augment_test.py +++ b/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() diff --git a/official/vision/image_classification/classifier_trainer.py b/official/vision/image_classification/classifier_trainer.py index 0c37f42378e..88ca5d147a4 100644 --- a/official/vision/image_classification/classifier_trainer.py +++ b/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) diff --git a/official/vision/image_classification/classifier_trainer_test.py b/official/vision/image_classification/classifier_trainer_test.py index 84c66c36daf..636650387d7 100644 --- a/official/vision/image_classification/classifier_trainer_test.py +++ b/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() diff --git a/official/vision/image_classification/dataset_factory.py b/official/vision/image_classification/dataset_factory.py index 27e770b40b4..703d892a02f 100644 --- a/official/vision/image_classification/dataset_factory.py +++ b/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 diff --git a/official/vision/image_classification/efficientnet/common_modules.py b/official/vision/image_classification/efficientnet/common_modules.py index b25dcb1c0e8..d37681eba01 100644 --- a/official/vision/image_classification/efficientnet/common_modules.py +++ b/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 diff --git a/official/vision/image_classification/efficientnet/efficientnet_model.py b/official/vision/image_classification/efficientnet/efficientnet_model.py index 657661b2c97..9a37d93ebf4 100644 --- a/official/vision/image_classification/efficientnet/efficientnet_model.py +++ b/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 diff --git a/official/vision/image_classification/learning_rate.py b/official/vision/image_classification/learning_rate.py index ae2b279866e..262b4aadc6e 100644 --- a/official/vision/image_classification/learning_rate.py +++ b/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 diff --git a/official/vision/image_classification/learning_rate_test.py b/official/vision/image_classification/learning_rate_test.py index 39b2bf11a3e..824115160ec 100644 --- a/official/vision/image_classification/learning_rate_test.py +++ b/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() diff --git a/official/vision/image_classification/mnist_test.py b/official/vision/image_classification/mnist_test.py index 0452e406c45..4c13db562d2 100644 --- a/official/vision/image_classification/mnist_test.py +++ b/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() diff --git a/official/vision/image_classification/optimizer_factory.py b/official/vision/image_classification/optimizer_factory.py index ae14b9f165d..14fcacfb410 100644 --- a/official/vision/image_classification/optimizer_factory.py +++ b/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 diff --git a/official/vision/image_classification/optimizer_factory_test.py b/official/vision/image_classification/optimizer_factory_test.py index 7d2a18ddb50..912732e2251 100644 --- a/official/vision/image_classification/optimizer_factory_test.py +++ b/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() diff --git a/official/vision/image_classification/preprocessing.py b/official/vision/image_classification/preprocessing.py index 200b228bfbf..3f2019189d4 100644 --- a/official/vision/image_classification/preprocessing.py +++ b/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 diff --git a/official/vision/image_classification/resnet/resnet_runnable.py b/official/vision/image_classification/resnet/resnet_runnable.py index 2bd1052e840..fd9816fc05c 100644 --- a/official/vision/image_classification/resnet/resnet_runnable.py +++ b/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 diff --git a/official/vision/image_classification/resnet/tfhub_export.py b/official/vision/image_classification/resnet/tfhub_export.py index 8ec000a220a..aefcc495eb5 100644 --- a/official/vision/image_classification/resnet/tfhub_export.py +++ b/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