Refactor tests for Object Detection API. (#8688)

Internal changes

--

PiperOrigin-RevId: 316837667

research/object_detection/models/ssd_mobilenet_v3_feature_extractor_test.py → research/object_detection/models/ssd_mobilenet_v3_feature_extractor_tf1_test.py

@@ -13,17 +13,15 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for ssd_mobilenet_v3_feature_extractor."""
-
+import unittest
 import tensorflow.compat.v1 as tf
-import tf_slim as slim
 
 from object_detection.models import ssd_mobilenet_v3_feature_extractor
 from object_detection.models import ssd_mobilenet_v3_feature_extractor_testbase
+from object_detection.utils import tf_version
 
 
-slim = slim
-
-
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SsdMobilenetV3LargeFeatureExtractorTest(
     ssd_mobilenet_v3_feature_extractor_testbase
     ._SsdMobilenetV3FeatureExtractorTestBase):
@@ -63,6 +61,7 @@ class SsdMobilenetV3LargeFeatureExtractorTest(
             use_explicit_padding=use_explicit_padding))
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SsdMobilenetV3SmallFeatureExtractorTest(
     ssd_mobilenet_v3_feature_extractor_testbase
     ._SsdMobilenetV3FeatureExtractorTestBase):

research/object_detection/models/ssd_pnasnet_feature_extractor.py

@@ -27,7 +27,10 @@ from object_detection.models import feature_map_generators
 from object_detection.utils import context_manager
 from object_detection.utils import ops
 from object_detection.utils import variables_helper
-from nets.nasnet import pnasnet
+try:
+  from nets.nasnet import pnasnet  # pylint: disable=g-import-not-at-top
+except:  # pylint: disable=bare-except
+  pass
 
 
 def pnasnet_large_arg_scope_for_detection(is_batch_norm_training=False):

research/object_detection/models/ssd_pnasnet_feature_extractor_test.py → research/object_detection/models/ssd_pnasnet_feature_extractor_tf1_test.py

@@ -14,13 +14,16 @@
 # ==============================================================================
 
 """Tests for ssd_pnas_feature_extractor."""
+import unittest
 import numpy as np
 import tensorflow.compat.v1 as tf
 
 from object_detection.models import ssd_feature_extractor_test
 from object_detection.models import ssd_pnasnet_feature_extractor
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SsdPnasNetFeatureExtractorTest(
     ssd_feature_extractor_test.SsdFeatureExtractorTestBase):
 

research/object_detection/models/ssd_resnet_v1_fpn_feature_extractor_testbase.py

@@ -19,24 +19,20 @@ from __future__ import division
 from __future__ import print_function
 
 import abc
-from absl.testing import parameterized
 import numpy as np
 from six.moves import zip
 import tensorflow.compat.v1 as tf
 
 from object_detection.models import ssd_feature_extractor_test
+from object_detection.utils import test_utils
 
 
-@parameterized.parameters(
-    {'use_keras': False},
-    {'use_keras': True},
-)
 class SSDResnetFPNFeatureExtractorTestBase(
     ssd_feature_extractor_test.SsdFeatureExtractorTestBase):
   """Helper test class for SSD Resnet v1 FPN feature extractors."""
 
   @abc.abstractmethod
-  def _resnet_scope_name(self, use_keras):
+  def _resnet_scope_name(self):
     pass
 
   @abc.abstractmethod
@@ -52,7 +48,7 @@ class SSDResnetFPNFeatureExtractorTestBase(
                                 use_keras=False):
     pass
 
-  def test_extract_features_returns_correct_shapes_256(self, use_keras):
+  def test_extract_features_returns_correct_shapes_256(self):
     image_height = 256
     image_width = 256
     depth_multiplier = 1.0
@@ -62,10 +58,10 @@ class SSDResnetFPNFeatureExtractorTestBase(
                                   (2, 2, 2, 256)]
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape)
+        expected_feature_map_shape, use_keras=self.is_tf2())
 
   def test_extract_features_returns_correct_shapes_with_dynamic_inputs(
-      self, use_keras):
+      self):
     image_height = 256
     image_width = 256
     depth_multiplier = 1.0
@@ -75,10 +71,10 @@ class SSDResnetFPNFeatureExtractorTestBase(
                                   (2, 2, 2, 256)]
     self.check_extract_features_returns_correct_shapes_with_dynamic_inputs(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape, use_keras=use_keras)
+        expected_feature_map_shape, use_keras=self.is_tf2())
 
   def test_extract_features_returns_correct_shapes_with_depth_multiplier(
-      self, use_keras):
+      self):
     image_height = 256
     image_width = 256
     depth_multiplier = 0.5
@@ -91,10 +87,10 @@ class SSDResnetFPNFeatureExtractorTestBase(
                                   (2, 2, 2, expected_num_channels)]
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape, use_keras=use_keras)
+        expected_feature_map_shape, use_keras=self.is_tf2())
 
   def test_extract_features_returns_correct_shapes_with_min_depth(
-      self, use_keras):
+      self):
     image_height = 256
     image_width = 256
     depth_multiplier = 1.0
@@ -106,23 +102,24 @@ class SSDResnetFPNFeatureExtractorTestBase(
                                   (2, 4, 4, min_depth),
                                   (2, 2, 2, min_depth)]
 
-    def graph_fn(image_tensor):
+    with test_utils.GraphContextOrNone() as g:
+      image_tensor = tf.random.uniform([2, image_height, image_width, 3])
       feature_extractor = self._create_feature_extractor(
           depth_multiplier, pad_to_multiple, min_depth=min_depth,
-          use_keras=use_keras)
-      if use_keras:
+          use_keras=self.is_tf2())
+
+    def graph_fn():
+      if self.is_tf2():
         return feature_extractor(image_tensor)
       return feature_extractor.extract_features(image_tensor)
 
-    image_tensor = np.random.rand(2, image_height, image_width,
-                                  3).astype(np.float32)
-    feature_maps = self.execute(graph_fn, [image_tensor])
+    feature_maps = self.execute(graph_fn, [], graph=g)
     for feature_map, expected_shape in zip(feature_maps,
                                            expected_feature_map_shape):
       self.assertAllEqual(feature_map.shape, expected_shape)
 
   def test_extract_features_returns_correct_shapes_with_pad_to_multiple(
-      self, use_keras):
+      self):
     image_height = 254
     image_width = 254
     depth_multiplier = 1.0
@@ -133,55 +130,62 @@ class SSDResnetFPNFeatureExtractorTestBase(
 
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape, use_keras=use_keras)
+        expected_feature_map_shape, use_keras=self.is_tf2())
 
   def test_extract_features_raises_error_with_invalid_image_size(
-      self, use_keras):
+      self):
     image_height = 32
     image_width = 32
     depth_multiplier = 1.0
     pad_to_multiple = 1
     self.check_extract_features_raises_error_with_invalid_image_size(
         image_height, image_width, depth_multiplier, pad_to_multiple,
-        use_keras=use_keras)
+        use_keras=self.is_tf2())
 
-  def test_preprocess_returns_correct_value_range(self, use_keras):
+  def test_preprocess_returns_correct_value_range(self):
     image_height = 128
     image_width = 128
     depth_multiplier = 1
     pad_to_multiple = 1
-    test_image = tf.constant(np.random.rand(4, image_height, image_width, 3))
-    feature_extractor = self._create_feature_extractor(depth_multiplier,
-                                                       pad_to_multiple,
-                                                       use_keras=use_keras)
-    preprocessed_image = feature_extractor.preprocess(test_image)
-    with self.test_session() as sess:
-      test_image_out, preprocessed_image_out = sess.run(
-          [test_image, preprocessed_image])
-      self.assertAllClose(preprocessed_image_out,
-                          test_image_out - [[123.68, 116.779, 103.939]])
-
-  def test_variables_only_created_in_scope(self, use_keras):
+    test_image_np = np.random.rand(4, image_height, image_width, 3)
+    with test_utils.GraphContextOrNone() as g:
+      test_image = tf.constant(test_image_np)
+      feature_extractor = self._create_feature_extractor(
+          depth_multiplier, pad_to_multiple, use_keras=self.is_tf2())
+
+    def graph_fn():
+      preprocessed_image = feature_extractor.preprocess(test_image)
+      return preprocessed_image
+
+    preprocessed_image_out = self.execute(graph_fn, [], graph=g)
+    self.assertAllClose(preprocessed_image_out,
+                        test_image_np - [[123.68, 116.779, 103.939]])
+
+  def test_variables_only_created_in_scope(self):
+    if self.is_tf2():
+      self.skipTest('test_variables_only_created_in_scope is only tf1')
     depth_multiplier = 1
     pad_to_multiple = 1
-    scope_name = self._resnet_scope_name(use_keras)
+    scope_name = self._resnet_scope_name()
     self.check_feature_extractor_variables_under_scope(
         depth_multiplier,
         pad_to_multiple,
         scope_name,
-        use_keras=use_keras)
+        use_keras=self.is_tf2())
 
-  def test_variable_count(self, use_keras):
+  def test_variable_count(self):
+    if self.is_tf2():
+      self.skipTest('test_variable_count is only tf1')
     depth_multiplier = 1
     pad_to_multiple = 1
     variables = self.get_feature_extractor_variables(
         depth_multiplier,
         pad_to_multiple,
-        use_keras=use_keras)
+        use_keras=self.is_tf2())
     # The number of expected variables in resnet_v1_50, resnet_v1_101,
     # and resnet_v1_152 is 279, 534, and 789 respectively.
     expected_variables_len = 279
-    scope_name = self._resnet_scope_name(use_keras)
+    scope_name = self._resnet_scope_name()
     if scope_name in ('ResNet101V1_FPN', 'resnet_v1_101'):
       expected_variables_len = 534
     elif scope_name in ('ResNet152V1_FPN', 'resnet_v1_152'):

research/object_detection/models/ssd_resnet_v1_fpn_feature_extractor_test.py → research/object_detection/models/ssd_resnet_v1_fpn_feature_extractor_tf1_test.py

@@ -13,13 +13,15 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for ssd resnet v1 FPN feature extractors."""
+import unittest
 import tensorflow.compat.v1 as tf
 
 from object_detection.models import ssd_resnet_v1_fpn_feature_extractor
 from object_detection.models import ssd_resnet_v1_fpn_feature_extractor_testbase
-from object_detection.models import ssd_resnet_v1_fpn_keras_feature_extractor
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SSDResnet50V1FeatureExtractorTest(
     ssd_resnet_v1_fpn_feature_extractor_testbase.
     SSDResnetFPNFeatureExtractorTestBase):
@@ -29,31 +31,17 @@ class SSDResnet50V1FeatureExtractorTest(
                                 use_explicit_padding=False, min_depth=32,
                                 use_keras=False):
     is_training = True
-    if use_keras:
-      return (ssd_resnet_v1_fpn_keras_feature_extractor.
-              SSDResNet50V1FpnKerasFeatureExtractor(
-                  is_training=is_training,
-                  depth_multiplier=depth_multiplier,
-                  min_depth=min_depth,
-                  pad_to_multiple=pad_to_multiple,
-                  conv_hyperparams=self._build_conv_hyperparams(
-                      add_batch_norm=False),
-                  freeze_batchnorm=False,
-                  inplace_batchnorm_update=False,
-                  name='ResNet50V1_FPN'))
-    else:
-      return (
-          ssd_resnet_v1_fpn_feature_extractor.SSDResnet50V1FpnFeatureExtractor(
-              is_training, depth_multiplier, min_depth, pad_to_multiple,
-              self.conv_hyperparams_fn,
-              use_explicit_padding=use_explicit_padding))
+    return (
+        ssd_resnet_v1_fpn_feature_extractor.SSDResnet50V1FpnFeatureExtractor(
+            is_training, depth_multiplier, min_depth, pad_to_multiple,
+            self.conv_hyperparams_fn,
+            use_explicit_padding=use_explicit_padding))
 
-  def _resnet_scope_name(self, use_keras=False):
-    if use_keras:
-      return 'ResNet50V1_FPN'
+  def _resnet_scope_name(self):
     return 'resnet_v1_50'
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SSDResnet101V1FeatureExtractorTest(
     ssd_resnet_v1_fpn_feature_extractor_testbase.
     SSDResnetFPNFeatureExtractorTestBase):
@@ -63,31 +51,17 @@ class SSDResnet101V1FeatureExtractorTest(
                                 use_explicit_padding=False, min_depth=32,
                                 use_keras=False):
     is_training = True
-    if use_keras:
-      return (ssd_resnet_v1_fpn_keras_feature_extractor.
-              SSDResNet101V1FpnKerasFeatureExtractor(
-                  is_training=is_training,
-                  depth_multiplier=depth_multiplier,
-                  min_depth=min_depth,
-                  pad_to_multiple=pad_to_multiple,
-                  conv_hyperparams=self._build_conv_hyperparams(
-                      add_batch_norm=False),
-                  freeze_batchnorm=False,
-                  inplace_batchnorm_update=False,
-                  name='ResNet101V1_FPN'))
-    else:
-      return (
-          ssd_resnet_v1_fpn_feature_extractor.SSDResnet101V1FpnFeatureExtractor(
-              is_training, depth_multiplier, min_depth, pad_to_multiple,
-              self.conv_hyperparams_fn,
-              use_explicit_padding=use_explicit_padding))
+    return (
+        ssd_resnet_v1_fpn_feature_extractor.SSDResnet101V1FpnFeatureExtractor(
+            is_training, depth_multiplier, min_depth, pad_to_multiple,
+            self.conv_hyperparams_fn,
+            use_explicit_padding=use_explicit_padding))
 
-  def _resnet_scope_name(self, use_keras):
-    if use_keras:
-      return 'ResNet101V1_FPN'
+  def _resnet_scope_name(self):
     return 'resnet_v1_101'
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SSDResnet152V1FeatureExtractorTest(
     ssd_resnet_v1_fpn_feature_extractor_testbase.
     SSDResnetFPNFeatureExtractorTestBase):
@@ -97,28 +71,13 @@ class SSDResnet152V1FeatureExtractorTest(
                                 use_explicit_padding=False, min_depth=32,
                                 use_keras=False):
     is_training = True
-    if use_keras:
-      return (ssd_resnet_v1_fpn_keras_feature_extractor.
-              SSDResNet152V1FpnKerasFeatureExtractor(
-                  is_training=is_training,
-                  depth_multiplier=depth_multiplier,
-                  min_depth=min_depth,
-                  pad_to_multiple=pad_to_multiple,
-                  conv_hyperparams=self._build_conv_hyperparams(
-                      add_batch_norm=False),
-                  freeze_batchnorm=False,
-                  inplace_batchnorm_update=False,
-                  name='ResNet152V1_FPN'))
-    else:
-      return (
-          ssd_resnet_v1_fpn_feature_extractor.SSDResnet152V1FpnFeatureExtractor(
-              is_training, depth_multiplier, min_depth, pad_to_multiple,
-              self.conv_hyperparams_fn,
-              use_explicit_padding=use_explicit_padding))
+    return (
+        ssd_resnet_v1_fpn_feature_extractor.SSDResnet152V1FpnFeatureExtractor(
+            is_training, depth_multiplier, min_depth, pad_to_multiple,
+            self.conv_hyperparams_fn,
+            use_explicit_padding=use_explicit_padding))
 
-  def _resnet_scope_name(self, use_keras):
-    if use_keras:
-      return 'ResNet152V1_FPN'
+  def _resnet_scope_name(self):
     return 'resnet_v1_152'
 
 

research/object_detection/models/ssd_resnet_v1_fpn_feature_extractor_tf2_test.py

+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ssd resnet v1 FPN feature extractors."""
+import unittest
+import tensorflow.compat.v1 as tf
+
+from object_detection.models import ssd_resnet_v1_fpn_feature_extractor_testbase
+from object_detection.models import ssd_resnet_v1_fpn_keras_feature_extractor
+from object_detection.utils import tf_version
+
+
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
+class SSDResnet50V1FeatureExtractorTest(
+    ssd_resnet_v1_fpn_feature_extractor_testbase.
+    SSDResnetFPNFeatureExtractorTestBase):
+  """SSDResnet50v1Fpn feature extractor test."""
+
+  def _create_feature_extractor(self, depth_multiplier, pad_to_multiple,
+                                use_explicit_padding=False, min_depth=32,
+                                use_keras=True):
+    is_training = True
+    return (ssd_resnet_v1_fpn_keras_feature_extractor.
+            SSDResNet50V1FpnKerasFeatureExtractor(
+                is_training=is_training,
+                depth_multiplier=depth_multiplier,
+                min_depth=min_depth,
+                pad_to_multiple=pad_to_multiple,
+                conv_hyperparams=self._build_conv_hyperparams(
+                    add_batch_norm=False),
+                freeze_batchnorm=False,
+                inplace_batchnorm_update=False,
+                name='ResNet50V1_FPN'))
+
+  def _resnet_scope_name(self):
+    return 'ResNet50V1_FPN'
+
+
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
+class SSDResnet101V1FeatureExtractorTest(
+    ssd_resnet_v1_fpn_feature_extractor_testbase.
+    SSDResnetFPNFeatureExtractorTestBase):
+  """SSDResnet101v1Fpn feature extractor test."""
+
+  def _create_feature_extractor(self, depth_multiplier, pad_to_multiple,
+                                use_explicit_padding=False, min_depth=32,
+                                use_keras=False):
+    is_training = True
+    return (ssd_resnet_v1_fpn_keras_feature_extractor.
+            SSDResNet101V1FpnKerasFeatureExtractor(
+                is_training=is_training,
+                depth_multiplier=depth_multiplier,
+                min_depth=min_depth,
+                pad_to_multiple=pad_to_multiple,
+                conv_hyperparams=self._build_conv_hyperparams(
+                    add_batch_norm=False),
+                freeze_batchnorm=False,
+                inplace_batchnorm_update=False,
+                name='ResNet101V1_FPN'))
+
+  def _resnet_scope_name(self):
+    return 'ResNet101V1_FPN'
+
+
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
+class SSDResnet152V1FeatureExtractorTest(
+    ssd_resnet_v1_fpn_feature_extractor_testbase.
+    SSDResnetFPNFeatureExtractorTestBase):
+  """SSDResnet152v1Fpn feature extractor test."""
+
+  def _create_feature_extractor(self, depth_multiplier, pad_to_multiple,
+                                use_explicit_padding=False, min_depth=32,
+                                use_keras=False):
+    is_training = True
+    return (ssd_resnet_v1_fpn_keras_feature_extractor.
+            SSDResNet152V1FpnKerasFeatureExtractor(
+                is_training=is_training,
+                depth_multiplier=depth_multiplier,
+                min_depth=min_depth,
+                pad_to_multiple=pad_to_multiple,
+                conv_hyperparams=self._build_conv_hyperparams(
+                    add_batch_norm=False),
+                freeze_batchnorm=False,
+                inplace_batchnorm_update=False,
+                name='ResNet152V1_FPN'))
+
+  def _resnet_scope_name(self):
+    return 'ResNet152V1_FPN'
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/object_detection/models/ssd_resnet_v1_ppn_feature_extractor_test.py → research/object_detection/models/ssd_resnet_v1_ppn_feature_extractor_tf1_test.py

@@ -13,12 +13,15 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for ssd resnet v1 feature extractors."""
+import unittest
 import tensorflow.compat.v1 as tf
 
 from object_detection.models import ssd_resnet_v1_ppn_feature_extractor
 from object_detection.models import ssd_resnet_v1_ppn_feature_extractor_testbase
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SSDResnet50V1PpnFeatureExtractorTest(
     ssd_resnet_v1_ppn_feature_extractor_testbase.
     SSDResnetPpnFeatureExtractorTestBase):
@@ -40,6 +43,7 @@ class SSDResnet50V1PpnFeatureExtractorTest(
     return 'resnet_v1_50'
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SSDResnet101V1PpnFeatureExtractorTest(
     ssd_resnet_v1_ppn_feature_extractor_testbase.
     SSDResnetPpnFeatureExtractorTestBase):
@@ -62,6 +66,7 @@ class SSDResnet101V1PpnFeatureExtractorTest(
     return 'resnet_v1_101'
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SSDResnet152V1PpnFeatureExtractorTest(
     ssd_resnet_v1_ppn_feature_extractor_testbase.
     SSDResnetPpnFeatureExtractorTestBase):

research/object_detection/predictors/convolutional_box_predictor_test.py → research/object_detection/predictors/convolutional_box_predictor_tf1_test.py

@@ -19,7 +19,7 @@
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
-
+import unittest
 from absl.testing import parameterized
 import numpy as np
 from six.moves import range
@@ -35,8 +35,10 @@ from object_detection.predictors.heads import class_head
 from object_detection.predictors.heads import mask_head
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class ConvolutionalBoxPredictorTest(test_case.TestCase):
 
   def _build_arg_scope_with_conv_hyperparams(self):
@@ -281,6 +283,7 @@ class ConvolutionalBoxPredictorTest(test_case.TestCase):
     self.assertEqual(bad_dangling_ops, [])
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class WeightSharedConvolutionalBoxPredictorTest(test_case.TestCase):
 
   def _build_arg_scope_with_conv_hyperparams(self):

research/object_detection/predictors/convolutional_keras_box_predictor_test.py → research/object_detection/predictors/convolutional_keras_box_predictor_tf2_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.convolutional_keras_box_predictor."""
+import unittest
 import numpy as np
 import tensorflow.compat.v1 as tf
 
@@ -26,8 +27,10 @@ from object_detection.predictors.heads import keras_class_head
 from object_detection.predictors.heads import keras_mask_head
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class ConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def _build_conv_hyperparams(self):
@@ -47,23 +50,23 @@ class ConvolutionalKerasBoxPredictorTest(test_case.TestCase):
     return hyperparams_builder.KerasLayerHyperparams(conv_hyperparams)
 
   def test_get_boxes_for_five_aspect_ratios_per_location(self):
+    conv_box_predictor = (
+        box_predictor_builder.build_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=0,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5],
+            min_depth=0,
+            max_depth=32,
+            num_layers_before_predictor=1,
+            use_dropout=True,
+            dropout_keep_prob=0.8,
+            kernel_size=1,
+            box_code_size=4
+        ))
     def graph_fn(image_features):
-      conv_box_predictor = (
-          box_predictor_builder.build_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=0,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5],
-              min_depth=0,
-              max_depth=32,
-              num_layers_before_predictor=1,
-              use_dropout=True,
-              dropout_keep_prob=0.8,
-              kernel_size=1,
-              box_code_size=4
-          ))
       box_predictions = conv_box_predictor([image_features])
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
@@ -78,23 +81,23 @@ class ConvolutionalKerasBoxPredictorTest(test_case.TestCase):
     self.assertAllEqual(objectness_predictions.shape, [4, 320, 1])
 
   def test_get_boxes_for_one_aspect_ratio_per_location(self):
+    conv_box_predictor = (
+        box_predictor_builder.build_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=0,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[1],
+            min_depth=0,
+            max_depth=32,
+            num_layers_before_predictor=1,
+            use_dropout=True,
+            dropout_keep_prob=0.8,
+            kernel_size=1,
+            box_code_size=4
+        ))
     def graph_fn(image_features):
-      conv_box_predictor = (
-          box_predictor_builder.build_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=0,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[1],
-              min_depth=0,
-              max_depth=32,
-              num_layers_before_predictor=1,
-              use_dropout=True,
-              dropout_keep_prob=0.8,
-              kernel_size=1,
-              box_code_size=4
-          ))
       box_predictions = conv_box_predictor([image_features])
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
@@ -111,23 +114,23 @@ class ConvolutionalKerasBoxPredictorTest(test_case.TestCase):
       self):
     num_classes_without_background = 6
     image_features = np.random.rand(4, 8, 8, 64).astype(np.float32)
+    conv_box_predictor = (
+        box_predictor_builder.build_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=num_classes_without_background,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5],
+            min_depth=0,
+            max_depth=32,
+            num_layers_before_predictor=1,
+            use_dropout=True,
+            dropout_keep_prob=0.8,
+            kernel_size=1,
+            box_code_size=4
+        ))
     def graph_fn(image_features):
-      conv_box_predictor = (
-          box_predictor_builder.build_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=num_classes_without_background,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5],
-              min_depth=0,
-              max_depth=32,
-              num_layers_before_predictor=1,
-              use_dropout=True,
-              dropout_keep_prob=0.8,
-              kernel_size=1,
-              box_code_size=4
-          ))
       box_predictions = conv_box_predictor([image_features])
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
@@ -144,7 +147,7 @@ class ConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def test_get_predictions_with_feature_maps_of_dynamic_shape(
       self):
-    image_features = tf.placeholder(dtype=tf.float32, shape=[4, None, None, 64])
+    tf.keras.backend.clear_session()
     conv_box_predictor = (
         box_predictor_builder.build_convolutional_keras_box_predictor(
             is_training=False,
@@ -161,28 +164,25 @@ class ConvolutionalKerasBoxPredictorTest(test_case.TestCase):
             kernel_size=1,
             box_code_size=4
         ))
-    box_predictions = conv_box_predictor([image_features])
-    box_encodings = tf.concat(
-        box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
-    objectness_predictions = tf.concat(
-        box_predictions[box_predictor.CLASS_PREDICTIONS_WITH_BACKGROUND],
-        axis=1)
-    init_op = tf.global_variables_initializer()
-
+    variables = []
+    def graph_fn(image_features):
+      box_predictions = conv_box_predictor([image_features])
+      variables.extend(list(conv_box_predictor.variables))
+      box_encodings = tf.concat(
+          box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
+      objectness_predictions = tf.concat(
+          box_predictions[box_predictor.CLASS_PREDICTIONS_WITH_BACKGROUND],
+          axis=1)
+      return box_encodings, objectness_predictions
     resolution = 32
     expected_num_anchors = resolution*resolution*5
-    with self.test_session() as sess:
-      sess.run(init_op)
-      (box_encodings_shape,
-       objectness_predictions_shape) = sess.run(
-           [tf.shape(box_encodings), tf.shape(objectness_predictions)],
-           feed_dict={image_features:
-                      np.random.rand(4, resolution, resolution, 64)})
-      actual_variable_set = set(
-          [var.op.name for var in tf.trainable_variables()])
-      self.assertAllEqual(box_encodings_shape, [4, expected_num_anchors, 1, 4])
-      self.assertAllEqual(objectness_predictions_shape,
-                          [4, expected_num_anchors, 1])
+    box_encodings, objectness_predictions = self.execute(
+        graph_fn, [np.random.rand(4, resolution, resolution, 64)])
+
+    actual_variable_set = set([var.name.split(':')[0] for var in variables])
+    self.assertAllEqual(box_encodings.shape, [4, expected_num_anchors, 1, 4])
+    self.assertAllEqual(objectness_predictions.shape,
+                        [4, expected_num_anchors, 1])
     expected_variable_set = set([
         'BoxPredictor/SharedConvolutions_0/Conv2d_0_1x1_32/bias',
         'BoxPredictor/SharedConvolutions_0/Conv2d_0_1x1_32/kernel',
@@ -195,7 +195,7 @@ class ConvolutionalKerasBoxPredictorTest(test_case.TestCase):
                      ['box_encodings', 'class_predictions_with_background'])
 
   def test_use_depthwise_convolution(self):
-    image_features = tf.placeholder(dtype=tf.float32, shape=[4, None, None, 64])
+    tf.keras.backend.clear_session()
     conv_box_predictor = (
         box_predictor_builder.build_convolutional_keras_box_predictor(
             is_training=False,
@@ -213,27 +213,25 @@ class ConvolutionalKerasBoxPredictorTest(test_case.TestCase):
             box_code_size=4,
             use_depthwise=True
         ))
-    box_predictions = conv_box_predictor([image_features])
-    box_encodings = tf.concat(
-        box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
-    objectness_predictions = tf.concat(
-        box_predictions[box_predictor.CLASS_PREDICTIONS_WITH_BACKGROUND],
-        axis=1)
-    init_op = tf.global_variables_initializer()
+    variables = []
+    def graph_fn(image_features):
+      box_predictions = conv_box_predictor([image_features])
+      variables.extend(list(conv_box_predictor.variables))
+      box_encodings = tf.concat(
+          box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
+      objectness_predictions = tf.concat(
+          box_predictions[box_predictor.CLASS_PREDICTIONS_WITH_BACKGROUND],
+          axis=1)
+      return box_encodings, objectness_predictions
 
     resolution = 32
     expected_num_anchors = resolution*resolution*5
-    with self.test_session() as sess:
-      sess.run(init_op)
-      (box_encodings_shape,
-       objectness_predictions_shape) = sess.run(
-           [tf.shape(box_encodings), tf.shape(objectness_predictions)],
-           feed_dict={image_features:
-                      np.random.rand(4, resolution, resolution, 64)})
-      actual_variable_set = set(
-          [var.op.name for var in tf.trainable_variables()])
-    self.assertAllEqual(box_encodings_shape, [4, expected_num_anchors, 1, 4])
-    self.assertAllEqual(objectness_predictions_shape,
+    box_encodings, objectness_predictions = self.execute(
+        graph_fn, [np.random.rand(4, resolution, resolution, 64)])
+
+    actual_variable_set = set([var.name.split(':')[0] for var in variables])
+    self.assertAllEqual(box_encodings.shape, [4, expected_num_anchors, 1, 4])
+    self.assertAllEqual(objectness_predictions.shape,
                         [4, expected_num_anchors, 1])
     expected_variable_set = set([
         'BoxPredictor/SharedConvolutions_0/Conv2d_0_1x1_32/bias',
@@ -259,6 +257,7 @@ class ConvolutionalKerasBoxPredictorTest(test_case.TestCase):
                      ['box_encodings', 'class_predictions_with_background'])
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def _build_conv_hyperparams(self, add_batch_norm=True):
@@ -288,19 +287,20 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   # pylint: disable=line-too-long
   def test_get_boxes_for_five_aspect_ratios_per_location(self):
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=0,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5],
+            depth=32,
+            num_layers_before_predictor=1,
+            box_code_size=4))
 
     def graph_fn(image_features):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=0,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5],
-              depth=32,
-              num_layers_before_predictor=1,
-              box_code_size=4))
       box_predictions = conv_box_predictor([image_features])
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
@@ -314,20 +314,21 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
     self.assertAllEqual(objectness_predictions.shape, [4, 320, 1])
 
   def test_bias_predictions_to_background_with_sigmoid_score_conversion(self):
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=True,
+            num_classes=2,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5],
+            depth=32,
+            num_layers_before_predictor=1,
+            class_prediction_bias_init=-4.6,
+            box_code_size=4))
 
     def graph_fn(image_features):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=True,
-              num_classes=2,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5],
-              depth=32,
-              num_layers_before_predictor=1,
-              class_prediction_bias_init=-4.6,
-              box_code_size=4))
       box_predictions = conv_box_predictor([image_features])
       class_predictions = tf.concat(box_predictions[
           box_predictor.CLASS_PREDICTIONS_WITH_BACKGROUND], axis=1)
@@ -339,20 +340,21 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def test_get_multi_class_predictions_for_five_aspect_ratios_per_location(
       self):
-
     num_classes_without_background = 6
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=num_classes_without_background,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5],
+            depth=32,
+            num_layers_before_predictor=1,
+            box_code_size=4))
+
     def graph_fn(image_features):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=num_classes_without_background,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5],
-              depth=32,
-              num_layers_before_predictor=1,
-              box_code_size=4))
       box_predictions = conv_box_predictor([image_features])
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
@@ -369,20 +371,21 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def test_get_multi_class_predictions_from_two_feature_maps(
       self):
-
     num_classes_without_background = 6
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=num_classes_without_background,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5, 5],
+            depth=32,
+            num_layers_before_predictor=1,
+            box_code_size=4))
+
     def graph_fn(image_features1, image_features2):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=num_classes_without_background,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5, 5],
-              depth=32,
-              num_layers_before_predictor=1,
-              box_code_size=4))
       box_predictions = conv_box_predictor([image_features1, image_features2])
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
@@ -401,20 +404,21 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def test_get_multi_class_predictions_from_feature_maps_of_different_depth(
       self):
-
     num_classes_without_background = 6
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=num_classes_without_background,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5, 5, 5],
+            depth=32,
+            num_layers_before_predictor=1,
+            box_code_size=4))
+
     def graph_fn(image_features1, image_features2, image_features3):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=num_classes_without_background,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5, 5, 5],
-              depth=32,
-              num_layers_before_predictor=1,
-              box_code_size=4))
       box_predictions = conv_box_predictor(
           [image_features1, image_features2, image_features3])
       box_encodings = tf.concat(
@@ -435,20 +439,25 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def test_predictions_multiple_feature_maps_share_weights_separate_batchnorm(
       self):
+    tf.keras.backend.clear_session()
     num_classes_without_background = 6
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=num_classes_without_background,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5, 5],
+            depth=32,
+            num_layers_before_predictor=2,
+            box_code_size=4))
+    variables = []
+
     def graph_fn(image_features1, image_features2):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=num_classes_without_background,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5, 5],
-              depth=32,
-              num_layers_before_predictor=2,
-              box_code_size=4))
       box_predictions = conv_box_predictor([image_features1, image_features2])
+      variables.extend(list(conv_box_predictor.variables))
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
       class_predictions_with_background = tf.concat(
@@ -456,25 +465,41 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
           axis=1)
       return (box_encodings, class_predictions_with_background)
 
-    with self.test_session(graph=tf.Graph()):
-      graph_fn(tf.random_uniform([4, 32, 32, 3], dtype=tf.float32),
-               tf.random_uniform([4, 16, 16, 3], dtype=tf.float32))
-      actual_variable_set = set(
-          [var.op.name for var in tf.trainable_variables()])
+    self.execute(graph_fn, [
+        np.random.rand(4, 32, 32, 3).astype(np.float32),
+        np.random.rand(4, 16, 16, 3).astype(np.float32)
+    ])
+    actual_variable_set = set([var.name.split(':')[0] for var in variables])
     expected_variable_set = set([
         # Box prediction tower
         ('WeightSharedConvolutionalBoxPredictor/'
          'BoxPredictionTower/conv2d_0/kernel'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'BoxPredictionTower/conv2d_0/BatchNorm/feature_0/beta'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'BoxPredictionTower/conv2d_0/BatchNorm/feature_0/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'BoxPredictionTower/conv2d_0/BatchNorm/feature_0/moving_variance'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'BoxPredictionTower/conv2d_0/BatchNorm/feature_1/beta'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'BoxPredictionTower/conv2d_0/BatchNorm/feature_1/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'BoxPredictionTower/conv2d_0/BatchNorm/feature_1/moving_variance'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'BoxPredictionTower/conv2d_1/kernel'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'BoxPredictionTower/conv2d_1/BatchNorm/feature_0/beta'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'BoxPredictionTower/conv2d_1/BatchNorm/feature_0/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'BoxPredictionTower/conv2d_1/BatchNorm/feature_0/moving_variance'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'BoxPredictionTower/conv2d_1/BatchNorm/feature_1/beta'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'BoxPredictionTower/conv2d_1/BatchNorm/feature_1/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'BoxPredictionTower/conv2d_1/BatchNorm/feature_1/moving_variance'),
         # Box prediction head
         ('WeightSharedConvolutionalBoxPredictor/'
          'WeightSharedConvolutionalBoxHead/BoxPredictor/kernel'),
@@ -485,14 +510,30 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
          'ClassPredictionTower/conv2d_0/kernel'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'ClassPredictionTower/conv2d_0/BatchNorm/feature_0/beta'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'ClassPredictionTower/conv2d_0/BatchNorm/feature_0/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'ClassPredictionTower/conv2d_0/BatchNorm/feature_0/moving_variance'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'ClassPredictionTower/conv2d_0/BatchNorm/feature_1/beta'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'ClassPredictionTower/conv2d_0/BatchNorm/feature_1/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'ClassPredictionTower/conv2d_0/BatchNorm/feature_1/moving_variance'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'ClassPredictionTower/conv2d_1/kernel'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'ClassPredictionTower/conv2d_1/BatchNorm/feature_0/beta'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'ClassPredictionTower/conv2d_1/BatchNorm/feature_0/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'ClassPredictionTower/conv2d_1/BatchNorm/feature_0/moving_variance'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'ClassPredictionTower/conv2d_1/BatchNorm/feature_1/beta'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'ClassPredictionTower/conv2d_1/BatchNorm/feature_1/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'ClassPredictionTower/conv2d_1/BatchNorm/feature_1/moving_variance'),
         # Class prediction head
         ('WeightSharedConvolutionalBoxPredictor/'
          'WeightSharedConvolutionalClassHead/ClassPredictor/kernel'),
@@ -502,21 +543,26 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def test_predictions_multiple_feature_maps_share_weights_without_batchnorm(
       self):
+    tf.keras.backend.clear_session()
     num_classes_without_background = 6
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=num_classes_without_background,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5, 5],
+            depth=32,
+            num_layers_before_predictor=2,
+            box_code_size=4,
+            apply_batch_norm=False))
+    variables = []
+
     def graph_fn(image_features1, image_features2):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=num_classes_without_background,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5, 5],
-              depth=32,
-              num_layers_before_predictor=2,
-              box_code_size=4,
-              apply_batch_norm=False))
       box_predictions = conv_box_predictor([image_features1, image_features2])
+      variables.extend(list(conv_box_predictor.variables))
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
       class_predictions_with_background = tf.concat(
@@ -524,11 +570,11 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
           axis=1)
       return (box_encodings, class_predictions_with_background)
 
-    with self.test_session(graph=tf.Graph()):
-      graph_fn(tf.random_uniform([4, 32, 32, 3], dtype=tf.float32),
-               tf.random_uniform([4, 16, 16, 3], dtype=tf.float32))
-      actual_variable_set = set(
-          [var.op.name for var in tf.trainable_variables()])
+    self.execute(graph_fn, [
+        np.random.rand(4, 32, 32, 3).astype(np.float32),
+        np.random.rand(4, 16, 16, 3).astype(np.float32)
+    ])
+    actual_variable_set = set([var.name.split(':')[0] for var in variables])
     expected_variable_set = set([
         # Box prediction tower
         ('WeightSharedConvolutionalBoxPredictor/'
@@ -562,23 +608,27 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def test_predictions_multiple_feature_maps_share_weights_with_depthwise(
       self):
+    tf.keras.backend.clear_session()
     num_classes_without_background = 6
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=num_classes_without_background,
+            conv_hyperparams=self._build_conv_hyperparams(add_batch_norm=False),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5, 5],
+            depth=32,
+            num_layers_before_predictor=2,
+            box_code_size=4,
+            apply_batch_norm=False,
+            use_depthwise=True))
+    variables = []
+
     def graph_fn(image_features1, image_features2):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=num_classes_without_background,
-              conv_hyperparams=self._build_conv_hyperparams(
-                  add_batch_norm=False),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5, 5],
-              depth=32,
-              num_layers_before_predictor=2,
-              box_code_size=4,
-              apply_batch_norm=False,
-              use_depthwise=True))
       box_predictions = conv_box_predictor([image_features1, image_features2])
+      variables.extend(list(conv_box_predictor.variables))
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
       class_predictions_with_background = tf.concat(
@@ -586,11 +636,11 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
           axis=1)
       return (box_encodings, class_predictions_with_background)
 
-    with self.test_session(graph=tf.Graph()):
-      graph_fn(tf.random_uniform([4, 32, 32, 3], dtype=tf.float32),
-               tf.random_uniform([4, 16, 16, 3], dtype=tf.float32))
-      actual_variable_set = set(
-          [var.op.name for var in tf.trainable_variables()])
+    self.execute(graph_fn, [
+        np.random.rand(4, 32, 32, 3).astype(np.float32),
+        np.random.rand(4, 16, 16, 3).astype(np.float32)
+    ])
+    actual_variable_set = set([var.name.split(':')[0] for var in variables])
     expected_variable_set = set([
         # Box prediction tower
         ('WeightSharedConvolutionalBoxPredictor/'
@@ -635,23 +685,27 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
     self.assertEqual(expected_variable_set, actual_variable_set)
 
   def test_no_batchnorm_params_when_batchnorm_is_not_configured(self):
+    tf.keras.backend.clear_session()
     num_classes_without_background = 6
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=num_classes_without_background,
+            conv_hyperparams=self._build_conv_hyperparams(add_batch_norm=False),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5, 5],
+            depth=32,
+            num_layers_before_predictor=2,
+            box_code_size=4,
+            apply_batch_norm=False))
+    variables = []
+
     def graph_fn(image_features1, image_features2):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=num_classes_without_background,
-              conv_hyperparams=self._build_conv_hyperparams(
-                  add_batch_norm=False),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5, 5],
-              depth=32,
-              num_layers_before_predictor=2,
-              box_code_size=4,
-              apply_batch_norm=False))
       box_predictions = conv_box_predictor(
           [image_features1, image_features2])
+      variables.extend(list(conv_box_predictor.variables))
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
       class_predictions_with_background = tf.concat(
@@ -659,11 +713,11 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
           axis=1)
       return (box_encodings, class_predictions_with_background)
 
-    with self.test_session(graph=tf.Graph()):
-      graph_fn(tf.random_uniform([4, 32, 32, 3], dtype=tf.float32),
-               tf.random_uniform([4, 16, 16, 3], dtype=tf.float32))
-      actual_variable_set = set(
-          [var.op.name for var in tf.trainable_variables()])
+    self.execute(graph_fn, [
+        np.random.rand(4, 32, 32, 3).astype(np.float32),
+        np.random.rand(4, 16, 16, 3).astype(np.float32)
+    ])
+    actual_variable_set = set([var.name.split(':')[0] for var in variables])
     expected_variable_set = set([
         # Box prediction tower
         ('WeightSharedConvolutionalBoxPredictor/'
@@ -697,22 +751,27 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def test_predictions_share_weights_share_tower_separate_batchnorm(
       self):
+    tf.keras.backend.clear_session()
     num_classes_without_background = 6
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=num_classes_without_background,
+            conv_hyperparams=self._build_conv_hyperparams(),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5, 5],
+            depth=32,
+            num_layers_before_predictor=2,
+            box_code_size=4,
+            share_prediction_tower=True))
+    variables = []
+
     def graph_fn(image_features1, image_features2):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=num_classes_without_background,
-              conv_hyperparams=self._build_conv_hyperparams(),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5, 5],
-              depth=32,
-              num_layers_before_predictor=2,
-              box_code_size=4,
-              share_prediction_tower=True))
       box_predictions = conv_box_predictor(
           [image_features1, image_features2])
+      variables.extend(list(conv_box_predictor.variables))
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
       class_predictions_with_background = tf.concat(
@@ -720,11 +779,11 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
           axis=1)
       return (box_encodings, class_predictions_with_background)
 
-    with self.test_session(graph=tf.Graph()):
-      graph_fn(tf.random_uniform([4, 32, 32, 3], dtype=tf.float32),
-               tf.random_uniform([4, 16, 16, 3], dtype=tf.float32))
-      actual_variable_set = set(
-          [var.op.name for var in tf.trainable_variables()])
+    self.execute(graph_fn, [
+        np.random.rand(4, 32, 32, 3).astype(np.float32),
+        np.random.rand(4, 16, 16, 3).astype(np.float32)
+    ])
+    actual_variable_set = set([var.name.split(':')[0] for var in variables])
     expected_variable_set = set([
         # Shared prediction tower
         ('WeightSharedConvolutionalBoxPredictor/'
@@ -733,12 +792,28 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
          'PredictionTower/conv2d_0/BatchNorm/feature_0/beta'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'PredictionTower/conv2d_0/BatchNorm/feature_1/beta'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'PredictionTower/conv2d_0/BatchNorm/feature_0/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'PredictionTower/conv2d_0/BatchNorm/feature_1/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'PredictionTower/conv2d_0/BatchNorm/feature_0/moving_variance'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'PredictionTower/conv2d_0/BatchNorm/feature_1/moving_variance'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'PredictionTower/conv2d_1/kernel'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'PredictionTower/conv2d_1/BatchNorm/feature_0/beta'),
         ('WeightSharedConvolutionalBoxPredictor/'
          'PredictionTower/conv2d_1/BatchNorm/feature_1/beta'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'PredictionTower/conv2d_1/BatchNorm/feature_0/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'PredictionTower/conv2d_1/BatchNorm/feature_1/moving_mean'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'PredictionTower/conv2d_1/BatchNorm/feature_0/moving_variance'),
+        ('WeightSharedConvolutionalBoxPredictor/'
+         'PredictionTower/conv2d_1/BatchNorm/feature_1/moving_variance'),
         # Box prediction head
         ('WeightSharedConvolutionalBoxPredictor/'
          'WeightSharedConvolutionalBoxHead/BoxPredictor/kernel'),
@@ -753,24 +828,28 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
   def test_predictions_share_weights_share_tower_without_batchnorm(
       self):
+    tf.keras.backend.clear_session()
     num_classes_without_background = 6
+    conv_box_predictor = (
+        box_predictor_builder
+        .build_weight_shared_convolutional_keras_box_predictor(
+            is_training=False,
+            num_classes=num_classes_without_background,
+            conv_hyperparams=self._build_conv_hyperparams(add_batch_norm=False),
+            freeze_batchnorm=False,
+            inplace_batchnorm_update=False,
+            num_predictions_per_location_list=[5, 5],
+            depth=32,
+            num_layers_before_predictor=2,
+            box_code_size=4,
+            share_prediction_tower=True,
+            apply_batch_norm=False))
+    variables = []
+
     def graph_fn(image_features1, image_features2):
-      conv_box_predictor = (
-          box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-              is_training=False,
-              num_classes=num_classes_without_background,
-              conv_hyperparams=self._build_conv_hyperparams(
-                  add_batch_norm=False),
-              freeze_batchnorm=False,
-              inplace_batchnorm_update=False,
-              num_predictions_per_location_list=[5, 5],
-              depth=32,
-              num_layers_before_predictor=2,
-              box_code_size=4,
-              share_prediction_tower=True,
-              apply_batch_norm=False))
       box_predictions = conv_box_predictor(
           [image_features1, image_features2])
+      variables.extend(list(conv_box_predictor.variables))
       box_encodings = tf.concat(
           box_predictions[box_predictor.BOX_ENCODINGS], axis=1)
       class_predictions_with_background = tf.concat(
@@ -778,11 +857,11 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
           axis=1)
       return (box_encodings, class_predictions_with_background)
 
-    with self.test_session(graph=tf.Graph()):
-      graph_fn(tf.random_uniform([4, 32, 32, 3], dtype=tf.float32),
-               tf.random_uniform([4, 16, 16, 3], dtype=tf.float32))
-      actual_variable_set = set(
-          [var.op.name for var in tf.trainable_variables()])
+    self.execute(graph_fn, [
+        np.random.rand(4, 32, 32, 3).astype(np.float32),
+        np.random.rand(4, 16, 16, 3).astype(np.float32)
+    ])
+    actual_variable_set = set([var.name.split(':')[0] for var in variables])
     expected_variable_set = set([
         # Shared prediction tower
         ('WeightSharedConvolutionalBoxPredictor/'
@@ -806,40 +885,6 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
 
     self.assertEqual(expected_variable_set, actual_variable_set)
 
-  def test_get_predictions_with_feature_maps_of_dynamic_shape(
-      self):
-    image_features = tf.placeholder(dtype=tf.float32, shape=[4, None, None, 64])
-    conv_box_predictor = (
-        box_predictor_builder.build_weight_shared_convolutional_keras_box_predictor(
-            is_training=False,
-            num_classes=0,
-            conv_hyperparams=self._build_conv_hyperparams(),
-            freeze_batchnorm=False,
-            inplace_batchnorm_update=False,
-            num_predictions_per_location_list=[5],
-            depth=32,
-            num_layers_before_predictor=1,
-            box_code_size=4))
-    box_predictions = conv_box_predictor([image_features])
-    box_encodings = tf.concat(box_predictions[box_predictor.BOX_ENCODINGS],
-                              axis=1)
-    objectness_predictions = tf.concat(box_predictions[
-        box_predictor.CLASS_PREDICTIONS_WITH_BACKGROUND], axis=1)
-    init_op = tf.global_variables_initializer()
-
-    resolution = 32
-    expected_num_anchors = resolution*resolution*5
-    with self.test_session() as sess:
-      sess.run(init_op)
-      (box_encodings_shape,
-       objectness_predictions_shape) = sess.run(
-           [tf.shape(box_encodings), tf.shape(objectness_predictions)],
-           feed_dict={image_features:
-                      np.random.rand(4, resolution, resolution, 64)})
-      self.assertAllEqual(box_encodings_shape, [4, expected_num_anchors, 4])
-      self.assertAllEqual(objectness_predictions_shape,
-                          [4, expected_num_anchors, 1])
-
   def test_other_heads_predictions(self):
     box_code_size = 4
     num_classes_without_background = 3
@@ -847,37 +892,36 @@ class WeightSharedConvolutionalKerasBoxPredictorTest(test_case.TestCase):
     mask_height = 5
     mask_width = 5
     num_predictions_per_location = 5
-
+    box_prediction_head = keras_box_head.WeightSharedConvolutionalBoxHead(
+        box_code_size=box_code_size,
+        conv_hyperparams=self._build_conv_hyperparams(),
+        num_predictions_per_location=num_predictions_per_location)
+    class_prediction_head = keras_class_head.WeightSharedConvolutionalClassHead(
+        num_class_slots=num_classes_without_background + 1,
+        conv_hyperparams=self._build_conv_hyperparams(),
+        num_predictions_per_location=num_predictions_per_location)
+    other_heads = {
+        other_head_name:
+            keras_mask_head.WeightSharedConvolutionalMaskHead(
+                num_classes=num_classes_without_background,
+                conv_hyperparams=self._build_conv_hyperparams(),
+                num_predictions_per_location=num_predictions_per_location,
+                mask_height=mask_height,
+                mask_width=mask_width)
+    }
+
+    conv_box_predictor = box_predictor.WeightSharedConvolutionalBoxPredictor(
+        is_training=False,
+        num_classes=num_classes_without_background,
+        box_prediction_head=box_prediction_head,
+        class_prediction_head=class_prediction_head,
+        other_heads=other_heads,
+        conv_hyperparams=self._build_conv_hyperparams(),
+        freeze_batchnorm=False,
+        inplace_batchnorm_update=False,
+        depth=32,
+        num_layers_before_predictor=2)
     def graph_fn(image_features):
-      box_prediction_head = keras_box_head.WeightSharedConvolutionalBoxHead(
-          box_code_size=box_code_size,
-          conv_hyperparams=self._build_conv_hyperparams(),
-          num_predictions_per_location=num_predictions_per_location)
-      class_prediction_head = keras_class_head.WeightSharedConvolutionalClassHead(
-          num_class_slots=num_classes_without_background + 1,
-          conv_hyperparams=self._build_conv_hyperparams(),
-          num_predictions_per_location=num_predictions_per_location)
-      other_heads = {
-          other_head_name:
-              keras_mask_head.WeightSharedConvolutionalMaskHead(
-                  num_classes=num_classes_without_background,
-                  conv_hyperparams=self._build_conv_hyperparams(),
-                  num_predictions_per_location=num_predictions_per_location,
-                  mask_height=mask_height,
-                  mask_width=mask_width)
-      }
-
-      conv_box_predictor = box_predictor.WeightSharedConvolutionalBoxPredictor(
-          is_training=False,
-          num_classes=num_classes_without_background,
-          box_prediction_head=box_prediction_head,
-          class_prediction_head=class_prediction_head,
-          other_heads=other_heads,
-          conv_hyperparams=self._build_conv_hyperparams(),
-          freeze_batchnorm=False,
-          inplace_batchnorm_update=False,
-          depth=32,
-          num_layers_before_predictor=2)
       box_predictions = conv_box_predictor([image_features])
       for key, value in box_predictions.items():
         box_predictions[key] = tf.concat(value, axis=1)

research/object_detection/predictors/heads/box_head_test.py → research/object_detection/predictors/heads/box_head_tf1_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.heads.box_head."""
+import unittest
 import tensorflow.compat.v1 as tf
 
 from google.protobuf import text_format
@@ -21,8 +22,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors.heads import box_head
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class MaskRCNNBoxHeadTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(self,
@@ -59,6 +62,7 @@ class MaskRCNNBoxHeadTest(test_case.TestCase):
     self.assertAllEqual([64, 1, 20, 4], prediction.get_shape().as_list())
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class ConvolutionalBoxPredictorTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(
@@ -92,6 +96,7 @@ class ConvolutionalBoxPredictorTest(test_case.TestCase):
     self.assertAllEqual([64, 323, 1, 4], box_encodings.get_shape().as_list())
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class WeightSharedConvolutionalBoxPredictorTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(

research/object_detection/predictors/heads/class_head_test.py → research/object_detection/predictors/heads/class_head_tf1_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.heads.class_head."""
+import unittest
 import tensorflow.compat.v1 as tf
 
 from google.protobuf import text_format
@@ -21,8 +22,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors.heads import class_head
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class MaskRCNNClassHeadTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(self,
@@ -81,6 +84,7 @@ class MaskRCNNClassHeadTest(test_case.TestCase):
       self.assertSetEqual(expected_var_names, actual_variable_set)
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class ConvolutionalClassPredictorTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(
@@ -140,6 +144,7 @@ class ConvolutionalClassPredictorTest(test_case.TestCase):
       self.assertSetEqual(expected_var_names, actual_variable_set)
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class WeightSharedConvolutionalClassPredictorTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(

research/object_detection/predictors/heads/keras_box_head_test.py → research/object_detection/predictors/heads/keras_box_head_tf2_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.heads.box_head."""
+import unittest
 import tensorflow.compat.v1 as tf
 
 from google.protobuf import text_format
@@ -21,8 +22,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors.heads import keras_box_head
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class ConvolutionalKerasBoxHeadTest(test_case.TestCase):
 
   def _build_conv_hyperparams(self):
@@ -51,10 +54,13 @@ class ConvolutionalKerasBoxHeadTest(test_case.TestCase):
         freeze_batchnorm=False,
         num_predictions_per_location=1,
         use_depthwise=False)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    box_encodings = box_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 1, 4], box_encodings.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      box_encodings = box_prediction_head(image_feature)
+      return box_encodings
+    box_encodings = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 1, 4], box_encodings.shape)
 
   def test_prediction_size_depthwise_true(self):
     conv_hyperparams = self._build_conv_hyperparams()
@@ -66,12 +72,16 @@ class ConvolutionalKerasBoxHeadTest(test_case.TestCase):
         freeze_batchnorm=False,
         num_predictions_per_location=1,
         use_depthwise=True)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    box_encodings = box_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 1, 4], box_encodings.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      box_encodings = box_prediction_head(image_feature)
+      return box_encodings
+    box_encodings = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 1, 4], box_encodings.shape)
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class MaskRCNNKerasBoxHeadTest(test_case.TestCase):
 
   def _build_fc_hyperparams(
@@ -102,12 +112,16 @@ class MaskRCNNKerasBoxHeadTest(test_case.TestCase):
         dropout_keep_prob=0.5,
         box_code_size=4,
         share_box_across_classes=False)
-    roi_pooled_features = tf.random_uniform(
-        [64, 7, 7, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    prediction = box_prediction_head(roi_pooled_features)
-    self.assertAllEqual([64, 1, 20, 4], prediction.get_shape().as_list())
+    def graph_fn():
+      roi_pooled_features = tf.random_uniform(
+          [64, 7, 7, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      prediction = box_prediction_head(roi_pooled_features)
+      return prediction
+    prediction = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 1, 20, 4], prediction.shape)
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class WeightSharedConvolutionalKerasBoxHead(test_case.TestCase):
 
   def _build_conv_hyperparams(self):
@@ -133,10 +147,13 @@ class WeightSharedConvolutionalKerasBoxHead(test_case.TestCase):
         conv_hyperparams=conv_hyperparams,
         num_predictions_per_location=1,
         use_depthwise=False)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    box_encodings = box_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 4], box_encodings.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      box_encodings = box_prediction_head(image_feature)
+      return box_encodings
+    box_encodings = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 4], box_encodings.shape)
 
   def test_prediction_size_depthwise_true(self):
     conv_hyperparams = self._build_conv_hyperparams()
@@ -145,40 +162,38 @@ class WeightSharedConvolutionalKerasBoxHead(test_case.TestCase):
         conv_hyperparams=conv_hyperparams,
         num_predictions_per_location=1,
         use_depthwise=True)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    box_encodings = box_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 4], box_encodings.get_shape().as_list())
-
-  def test_variable_count_depth_wise_true(self):
-    g = tf.Graph()
-    with g.as_default():
-      conv_hyperparams = self._build_conv_hyperparams()
-      box_prediction_head = keras_box_head.WeightSharedConvolutionalBoxHead(
-          box_code_size=4,
-          conv_hyperparams=conv_hyperparams,
-          num_predictions_per_location=1,
-          use_depthwise=True)
+    def graph_fn():
       image_feature = tf.random_uniform(
           [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-      _ = box_prediction_head(image_feature)
-      variables = g.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
-    self.assertEqual(len(variables), 3)
+      box_encodings = box_prediction_head(image_feature)
+      return box_encodings
+    box_encodings = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 4], box_encodings.shape)
+
+  def test_variable_count_depth_wise_true(self):
+    conv_hyperparams = self._build_conv_hyperparams()
+    box_prediction_head = keras_box_head.WeightSharedConvolutionalBoxHead(
+        box_code_size=4,
+        conv_hyperparams=conv_hyperparams,
+        num_predictions_per_location=1,
+        use_depthwise=True)
+    image_feature = tf.random_uniform(
+        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+    box_prediction_head(image_feature)
+    self.assertEqual(len(box_prediction_head.variables), 3)
 
   def test_variable_count_depth_wise_False(self):
-    g = tf.Graph()
-    with g.as_default():
-      conv_hyperparams = self._build_conv_hyperparams()
-      box_prediction_head = keras_box_head.WeightSharedConvolutionalBoxHead(
-          box_code_size=4,
-          conv_hyperparams=conv_hyperparams,
-          num_predictions_per_location=1,
-          use_depthwise=False)
-      image_feature = tf.random_uniform(
-          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-      _ = box_prediction_head(image_feature)
-      variables = g.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
-    self.assertEqual(len(variables), 2)
+    conv_hyperparams = self._build_conv_hyperparams()
+    box_prediction_head = keras_box_head.WeightSharedConvolutionalBoxHead(
+        box_code_size=4,
+        conv_hyperparams=conv_hyperparams,
+        num_predictions_per_location=1,
+        use_depthwise=False)
+    image_feature = tf.random_uniform(
+        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+    box_prediction_head(image_feature)
+    self.assertEqual(len(box_prediction_head.variables), 2)
+
 
 if __name__ == '__main__':
   tf.test.main()

research/object_detection/predictors/heads/keras_class_head_test.py → research/object_detection/predictors/heads/keras_class_head_tf2_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.heads.class_head."""
+import unittest
 import tensorflow.compat.v1 as tf
 
 from google.protobuf import text_format
@@ -21,8 +22,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors.heads import keras_class_head
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class ConvolutionalKerasClassPredictorTest(test_case.TestCase):
 
   def _build_conv_hyperparams(self):
@@ -53,11 +56,13 @@ class ConvolutionalKerasClassPredictorTest(test_case.TestCase):
         freeze_batchnorm=False,
         num_predictions_per_location=1,
         use_depthwise=False)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    class_predictions = class_prediction_head(image_feature,)
-    self.assertAllEqual([64, 323, 20],
-                        class_predictions.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      class_predictions = class_prediction_head(image_feature,)
+      return class_predictions
+    class_predictions = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 20], class_predictions.shape)
 
   def test_prediction_size_depthwise_true(self):
     conv_hyperparams = self._build_conv_hyperparams()
@@ -71,13 +76,16 @@ class ConvolutionalKerasClassPredictorTest(test_case.TestCase):
         freeze_batchnorm=False,
         num_predictions_per_location=1,
         use_depthwise=True)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    class_predictions = class_prediction_head(image_feature,)
-    self.assertAllEqual([64, 323, 20],
-                        class_predictions.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      class_predictions = class_prediction_head(image_feature,)
+      return class_predictions
+    class_predictions = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 20], class_predictions.shape)
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class MaskRCNNClassHeadTest(test_case.TestCase):
 
   def _build_fc_hyperparams(self,
@@ -106,12 +114,16 @@ class MaskRCNNClassHeadTest(test_case.TestCase):
         freeze_batchnorm=False,
         use_dropout=True,
         dropout_keep_prob=0.5)
-    roi_pooled_features = tf.random_uniform(
-        [64, 7, 7, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    prediction = class_prediction_head(roi_pooled_features)
-    self.assertAllEqual([64, 1, 20], prediction.get_shape().as_list())
+    def graph_fn():
+      roi_pooled_features = tf.random_uniform(
+          [64, 7, 7, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      prediction = class_prediction_head(roi_pooled_features)
+      return prediction
+    prediction = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 1, 20], prediction.shape)
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class WeightSharedConvolutionalKerasClassPredictorTest(test_case.TestCase):
 
   def _build_conv_hyperparams(self):
@@ -137,10 +149,13 @@ class WeightSharedConvolutionalKerasClassPredictorTest(test_case.TestCase):
         conv_hyperparams=conv_hyperparams,
         num_predictions_per_location=1,
         use_depthwise=False)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    class_predictions = class_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 20], class_predictions.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      class_predictions = class_prediction_head(image_feature)
+      return class_predictions
+    class_predictions = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 20], class_predictions.shape)
 
   def test_prediction_size_depthwise_true(self):
     conv_hyperparams = self._build_conv_hyperparams()
@@ -149,42 +164,39 @@ class WeightSharedConvolutionalKerasClassPredictorTest(test_case.TestCase):
         conv_hyperparams=conv_hyperparams,
         num_predictions_per_location=1,
         use_depthwise=True)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    class_predictions = class_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 20], class_predictions.get_shape().as_list())
-
-  def test_variable_count_depth_wise_true(self):
-    g = tf.Graph()
-    with g.as_default():
-      conv_hyperparams = self._build_conv_hyperparams()
-      class_prediction_head = (
-          keras_class_head.WeightSharedConvolutionalClassHead(
-              num_class_slots=20,
-              conv_hyperparams=conv_hyperparams,
-              num_predictions_per_location=1,
-              use_depthwise=True))
+    def graph_fn():
       image_feature = tf.random_uniform(
           [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-      _ = class_prediction_head(image_feature)
-      variables = g.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
-    self.assertEqual(len(variables), 3)
+      class_predictions = class_prediction_head(image_feature)
+      return class_predictions
+    class_predictions = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 20], class_predictions.shape)
+
+  def test_variable_count_depth_wise_true(self):
+    conv_hyperparams = self._build_conv_hyperparams()
+    class_prediction_head = (
+        keras_class_head.WeightSharedConvolutionalClassHead(
+            num_class_slots=20,
+            conv_hyperparams=conv_hyperparams,
+            num_predictions_per_location=1,
+            use_depthwise=True))
+    image_feature = tf.random_uniform(
+        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+    class_prediction_head(image_feature)
+    self.assertEqual(len(class_prediction_head.variables), 3)
 
   def test_variable_count_depth_wise_False(self):
-    g = tf.Graph()
-    with g.as_default():
-      conv_hyperparams = self._build_conv_hyperparams()
-      class_prediction_head = (
-          keras_class_head.WeightSharedConvolutionalClassHead(
-              num_class_slots=20,
-              conv_hyperparams=conv_hyperparams,
-              num_predictions_per_location=1,
-              use_depthwise=False))
-      image_feature = tf.random_uniform(
-          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-      _ = class_prediction_head(image_feature)
-      variables = g.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
-    self.assertEqual(len(variables), 2)
+    conv_hyperparams = self._build_conv_hyperparams()
+    class_prediction_head = (
+        keras_class_head.WeightSharedConvolutionalClassHead(
+            num_class_slots=20,
+            conv_hyperparams=conv_hyperparams,
+            num_predictions_per_location=1,
+            use_depthwise=False))
+    image_feature = tf.random_uniform(
+        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+    class_prediction_head(image_feature)
+    self.assertEqual(len(class_prediction_head.variables), 2)
 
 
 if __name__ == '__main__':

research/object_detection/predictors/heads/keras_mask_head_test.py → research/object_detection/predictors/heads/keras_mask_head_tf2_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.heads.mask_head."""
+import unittest
 import tensorflow.compat.v1 as tf
 
 from google.protobuf import text_format
@@ -21,8 +22,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors.heads import keras_mask_head
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class ConvolutionalMaskPredictorTest(test_case.TestCase):
 
   def _build_conv_hyperparams(self):
@@ -55,11 +58,13 @@ class ConvolutionalMaskPredictorTest(test_case.TestCase):
         use_depthwise=False,
         mask_height=7,
         mask_width=7)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    mask_predictions = mask_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 20, 7, 7],
-                        mask_predictions.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      mask_predictions = mask_prediction_head(image_feature)
+      return mask_predictions
+    mask_predictions = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 20, 7, 7], mask_predictions.shape)
 
   def test_prediction_size_use_depthwise_true(self):
     conv_hyperparams = self._build_conv_hyperparams()
@@ -75,11 +80,13 @@ class ConvolutionalMaskPredictorTest(test_case.TestCase):
         use_depthwise=True,
         mask_height=7,
         mask_width=7)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    mask_predictions = mask_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 20, 7, 7],
-                        mask_predictions.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      mask_predictions = mask_prediction_head(image_feature)
+      return mask_predictions
+    mask_predictions = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 20, 7, 7], mask_predictions.shape)
 
   def test_class_agnostic_prediction_size_use_depthwise_false(self):
     conv_hyperparams = self._build_conv_hyperparams()
@@ -96,11 +103,13 @@ class ConvolutionalMaskPredictorTest(test_case.TestCase):
         mask_height=7,
         mask_width=7,
         masks_are_class_agnostic=True)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    mask_predictions = mask_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 1, 7, 7],
-                        mask_predictions.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      mask_predictions = mask_prediction_head(image_feature)
+      return mask_predictions
+    mask_predictions = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 1, 7, 7], mask_predictions.shape)
 
   def test_class_agnostic_prediction_size_use_depthwise_true(self):
     conv_hyperparams = self._build_conv_hyperparams()
@@ -117,13 +126,16 @@ class ConvolutionalMaskPredictorTest(test_case.TestCase):
         mask_height=7,
         mask_width=7,
         masks_are_class_agnostic=True)
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    mask_predictions = mask_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 1, 7, 7],
-                        mask_predictions.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      mask_predictions = mask_prediction_head(image_feature)
+      return mask_predictions
+    mask_predictions = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 1, 7, 7], mask_predictions.shape)
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class MaskRCNNMaskHeadTest(test_case.TestCase):
 
   def _build_conv_hyperparams(self,
@@ -155,10 +167,13 @@ class MaskRCNNMaskHeadTest(test_case.TestCase):
         mask_prediction_num_conv_layers=2,
         mask_prediction_conv_depth=256,
         masks_are_class_agnostic=False)
-    roi_pooled_features = tf.random_uniform(
-        [64, 7, 7, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    prediction = mask_prediction_head(roi_pooled_features)
-    self.assertAllEqual([64, 1, 20, 14, 14], prediction.get_shape().as_list())
+    def graph_fn():
+      roi_pooled_features = tf.random_uniform(
+          [64, 7, 7, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      prediction = mask_prediction_head(roi_pooled_features)
+      return prediction
+    prediction = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 1, 20, 14, 14], prediction.shape)
 
   def test_prediction_size_with_convolve_then_upsample(self):
     mask_prediction_head = keras_mask_head.MaskRCNNMaskHead(
@@ -172,12 +187,16 @@ class MaskRCNNMaskHeadTest(test_case.TestCase):
         mask_prediction_conv_depth=256,
         masks_are_class_agnostic=True,
         convolve_then_upsample=True)
-    roi_pooled_features = tf.random_uniform(
-        [64, 14, 14, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    prediction = mask_prediction_head(roi_pooled_features)
-    self.assertAllEqual([64, 1, 1, 28, 28], prediction.get_shape().as_list())
+    def graph_fn():
+      roi_pooled_features = tf.random_uniform(
+          [64, 14, 14, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      prediction = mask_prediction_head(roi_pooled_features)
+      return prediction
+    prediction = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 1, 1, 28, 28], prediction.shape)
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class WeightSharedConvolutionalMaskPredictorTest(test_case.TestCase):
 
   def _build_conv_hyperparams(self):
@@ -204,11 +223,13 @@ class WeightSharedConvolutionalMaskPredictorTest(test_case.TestCase):
             conv_hyperparams=self._build_conv_hyperparams(),
             mask_height=7,
             mask_width=7))
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    mask_predictions = mask_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 20, 7, 7],
-                        mask_predictions.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      mask_predictions = mask_prediction_head(image_feature)
+      return mask_predictions
+    mask_predictions = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 20, 7, 7], mask_predictions.shape)
 
   def test_class_agnostic_prediction_size(self):
     mask_prediction_head = (
@@ -219,11 +240,13 @@ class WeightSharedConvolutionalMaskPredictorTest(test_case.TestCase):
             mask_height=7,
             mask_width=7,
             masks_are_class_agnostic=True))
-    image_feature = tf.random_uniform(
-        [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
-    mask_predictions = mask_prediction_head(image_feature)
-    self.assertAllEqual([64, 323, 1, 7, 7],
-                        mask_predictions.get_shape().as_list())
+    def graph_fn():
+      image_feature = tf.random_uniform(
+          [64, 17, 19, 1024], minval=-10.0, maxval=10.0, dtype=tf.float32)
+      mask_predictions = mask_prediction_head(image_feature)
+      return mask_predictions
+    mask_predictions = self.execute(graph_fn, [])
+    self.assertAllEqual([64, 323, 1, 7, 7], mask_predictions.shape)
 
 if __name__ == '__main__':
   tf.test.main()

research/object_detection/predictors/heads/keypoint_head_test.py → research/object_detection/predictors/heads/keypoint_head_tf1_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.heads.keypoint_head."""
+import unittest
 import tensorflow.compat.v1 as tf
 
 from google.protobuf import text_format
@@ -21,8 +22,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors.heads import keypoint_head
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class MaskRCNNKeypointHeadTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(self,

research/object_detection/predictors/heads/mask_head_test.py → research/object_detection/predictors/heads/mask_head_tf1_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.heads.mask_head."""
+import unittest
 import tensorflow.compat.v1 as tf
 
 from google.protobuf import text_format
@@ -21,8 +22,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors.heads import mask_head
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class MaskRCNNMaskHeadTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(self,
@@ -75,6 +78,7 @@ class MaskRCNNMaskHeadTest(test_case.TestCase):
     self.assertAllEqual([64, 1, 1, 28, 28], prediction.get_shape().as_list())
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class ConvolutionalMaskPredictorTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(
@@ -131,6 +135,7 @@ class ConvolutionalMaskPredictorTest(test_case.TestCase):
                         mask_predictions.get_shape().as_list())
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class WeightSharedConvolutionalMaskPredictorTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(

research/object_detection/predictors/mask_rcnn_box_predictor_test.py → research/object_detection/predictors/mask_rcnn_box_predictor_tf1_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.mask_rcnn_box_predictor."""
+import unittest
 import numpy as np
 import tensorflow.compat.v1 as tf
 
@@ -23,8 +24,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors import mask_rcnn_box_predictor as box_predictor
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class MaskRCNNBoxPredictorTest(test_case.TestCase):
 
   def _build_arg_scope_with_hyperparams(self,

research/object_detection/predictors/mask_rcnn_keras_box_predictor_test.py → research/object_detection/predictors/mask_rcnn_keras_box_predictor_tf2_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.mask_rcnn_box_predictor."""
+import unittest
 import numpy as np
 import tensorflow.compat.v1 as tf
 
@@ -23,8 +24,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors import mask_rcnn_keras_box_predictor as box_predictor
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class MaskRCNNKerasBoxPredictorTest(test_case.TestCase):
 
   def _build_hyperparams(self,
@@ -46,17 +49,17 @@ class MaskRCNNKerasBoxPredictorTest(test_case.TestCase):
     return hyperparams_builder.KerasLayerHyperparams(hyperparams)
 
   def test_get_boxes_with_five_classes(self):
+    mask_box_predictor = (
+        box_predictor_builder.build_mask_rcnn_keras_box_predictor(
+            is_training=False,
+            num_classes=5,
+            fc_hyperparams=self._build_hyperparams(),
+            freeze_batchnorm=False,
+            use_dropout=False,
+            dropout_keep_prob=0.5,
+            box_code_size=4,
+        ))
     def graph_fn(image_features):
-      mask_box_predictor = (
-          box_predictor_builder.build_mask_rcnn_keras_box_predictor(
-              is_training=False,
-              num_classes=5,
-              fc_hyperparams=self._build_hyperparams(),
-              freeze_batchnorm=False,
-              use_dropout=False,
-              dropout_keep_prob=0.5,
-              box_code_size=4,
-          ))
       box_predictions = mask_box_predictor(
           [image_features],
           prediction_stage=2)
@@ -70,18 +73,19 @@ class MaskRCNNKerasBoxPredictorTest(test_case.TestCase):
     self.assertAllEqual(class_predictions_with_background.shape, [2, 1, 6])
 
   def test_get_boxes_with_five_classes_share_box_across_classes(self):
+    mask_box_predictor = (
+        box_predictor_builder.build_mask_rcnn_keras_box_predictor(
+            is_training=False,
+            num_classes=5,
+            fc_hyperparams=self._build_hyperparams(),
+            freeze_batchnorm=False,
+            use_dropout=False,
+            dropout_keep_prob=0.5,
+            box_code_size=4,
+            share_box_across_classes=True
+        ))
     def graph_fn(image_features):
-      mask_box_predictor = (
-          box_predictor_builder.build_mask_rcnn_keras_box_predictor(
-              is_training=False,
-              num_classes=5,
-              fc_hyperparams=self._build_hyperparams(),
-              freeze_batchnorm=False,
-              use_dropout=False,
-              dropout_keep_prob=0.5,
-              box_code_size=4,
-              share_box_across_classes=True
-          ))
+
       box_predictions = mask_box_predictor(
           [image_features],
           prediction_stage=2)
@@ -95,19 +99,19 @@ class MaskRCNNKerasBoxPredictorTest(test_case.TestCase):
     self.assertAllEqual(class_predictions_with_background.shape, [2, 1, 6])
 
   def test_get_instance_masks(self):
+    mask_box_predictor = (
+        box_predictor_builder.build_mask_rcnn_keras_box_predictor(
+            is_training=False,
+            num_classes=5,
+            fc_hyperparams=self._build_hyperparams(),
+            freeze_batchnorm=False,
+            use_dropout=False,
+            dropout_keep_prob=0.5,
+            box_code_size=4,
+            conv_hyperparams=self._build_hyperparams(
+                op_type=hyperparams_pb2.Hyperparams.CONV),
+            predict_instance_masks=True))
     def graph_fn(image_features):
-      mask_box_predictor = (
-          box_predictor_builder.build_mask_rcnn_keras_box_predictor(
-              is_training=False,
-              num_classes=5,
-              fc_hyperparams=self._build_hyperparams(),
-              freeze_batchnorm=False,
-              use_dropout=False,
-              dropout_keep_prob=0.5,
-              box_code_size=4,
-              conv_hyperparams=self._build_hyperparams(
-                  op_type=hyperparams_pb2.Hyperparams.CONV),
-              predict_instance_masks=True))
       box_predictions = mask_box_predictor(
           [image_features],
           prediction_stage=3)

research/object_detection/predictors/rfcn_box_predictor_test.py → research/object_detection/predictors/rfcn_box_predictor_tf1_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.rfcn_box_predictor."""
+import unittest
 import numpy as np
 import tensorflow.compat.v1 as tf
 
@@ -22,8 +23,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors import rfcn_box_predictor as box_predictor
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class RfcnBoxPredictorTest(test_case.TestCase):
 
   def _build_arg_scope_with_conv_hyperparams(self):

research/object_detection/predictors/rfcn_keras_box_predictor_test.py → research/object_detection/predictors/rfcn_keras_box_predictor_tf2_test.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 """Tests for object_detection.predictors.rfcn_box_predictor."""
+import unittest
 import numpy as np
 import tensorflow.compat.v1 as tf
 
@@ -22,8 +23,10 @@ from object_detection.builders import hyperparams_builder
 from object_detection.predictors import rfcn_keras_box_predictor as box_predictor
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
 class RfcnKerasBoxPredictorTest(test_case.TestCase):
 
   def _build_conv_hyperparams(self):
@@ -42,18 +45,17 @@ class RfcnKerasBoxPredictorTest(test_case.TestCase):
     return hyperparams_builder.KerasLayerHyperparams(conv_hyperparams)
 
   def test_get_correct_box_encoding_and_class_prediction_shapes(self):
-
+    rfcn_box_predictor = box_predictor.RfcnKerasBoxPredictor(
+        is_training=False,
+        num_classes=2,
+        conv_hyperparams=self._build_conv_hyperparams(),
+        freeze_batchnorm=False,
+        num_spatial_bins=[3, 3],
+        depth=4,
+        crop_size=[12, 12],
+        box_code_size=4)
     def graph_fn(image_features, proposal_boxes):
-      rfcn_box_predictor = box_predictor.RfcnKerasBoxPredictor(
-          is_training=False,
-          num_classes=2,
-          conv_hyperparams=self._build_conv_hyperparams(),
-          freeze_batchnorm=False,
-          num_spatial_bins=[3, 3],
-          depth=4,
-          crop_size=[12, 12],
-          box_code_size=4
-      )
+
       box_predictions = rfcn_box_predictor(
           [image_features],
           proposal_boxes=proposal_boxes)