Merge remote-tracking branch 'upstream/master' into fasterrcnn_fpn_keras_feature_extractor

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)

research/object_detection/protos/center_net.proto

+syntax = "proto2";
+
+package object_detection.protos;
+
+import "object_detection/protos/image_resizer.proto";
+import "object_detection/protos/losses.proto";
+
+// Configuration for the CenterNet meta architecture from the "Objects as
+// Points" paper [1]
+// [1]: https://arxiv.org/abs/1904.07850
+
+message CenterNet {
+  // Number of classes to predict.
+  optional int32 num_classes = 1;
+
+  // Feature extractor config.
+  optional CenterNetFeatureExtractor feature_extractor = 2;
+
+  // Image resizer for preprocessing the input image.
+  optional ImageResizer image_resizer = 3;
+
+  // Parameters which are related to object detection task.
+  message ObjectDetection {
+    // The original fields are moved to ObjectCenterParams or deleted.
+    reserved 2, 5, 6, 7;
+
+    // Weight of the task loss. The total loss of the model will be the
+    // summation of task losses weighted by the weights.
+    optional float task_loss_weight = 1 [default = 1.0];
+
+    // Weight for the offset localization loss.
+    optional float offset_loss_weight = 3 [default = 1.0];
+
+    // Weight for the height/width localization loss.
+    optional float scale_loss_weight = 4 [default = 0.1];
+
+    // Localization loss configuration for object scale and offset losses.
+    optional LocalizationLoss localization_loss = 8;
+  }
+  optional ObjectDetection object_detection_task = 4;
+
+  // Parameters related to object center prediction. This is required for both
+  // object detection and keypoint estimation tasks.
+  message ObjectCenterParams {
+    // Weight for the object center loss.
+    optional float object_center_loss_weight = 1 [default = 1.0];
+
+    // Classification loss configuration for object center loss.
+    optional ClassificationLoss classification_loss = 2;
+
+    // The initial bias value of the convlution kernel of the class heatmap
+    // prediction head. -2.19 corresponds to predicting foreground with
+    // a probability of 0.1. See "Focal Loss for Dense Object Detection"
+    // at https://arxiv.org/abs/1708.02002.
+    optional float heatmap_bias_init = 3 [default = -2.19];
+
+    // The minimum IOU overlap boxes need to have to not be penalized.
+    optional float min_box_overlap_iou = 4 [default = 0.7];
+
+    // Maximum number of boxes to predict.
+    optional int32 max_box_predictions = 5 [default = 100];
+
+    // If set, loss is only computed for the labeled classes.
+    optional bool use_labeled_classes = 6 [default = false];
+  }
+  optional ObjectCenterParams object_center_params = 5;
+
+  // Path of the file that conatins the label map along with the keypoint
+  // information, including the keypoint indices, corresponding labels, and the
+  // corresponding class. The file should be the same one as used in the input
+  // pipeline. Note that a plain text of StringIntLabelMap proto is expected in
+  // this file.
+  // It is required only if the keypoint estimation task is specified.
+  optional string keypoint_label_map_path = 6;
+
+  // Parameters which are related to keypoint estimation task.
+  message KeypointEstimation {
+    // Name of the task, e.g. "human pose". Note that the task name should be
+    // unique to each keypoint task.
+    optional string task_name = 1;
+
+    // Weight of the task loss. The total loss of the model will be their
+    // summation of task losses weighted by the weights.
+    optional float task_loss_weight = 2 [default = 1.0];
+
+    // Loss configuration for keypoint heatmap, offset, regression losses. Note
+    // that the localization loss is used for offset/regression losses and
+    // classification loss is used for heatmap loss.
+    optional Loss loss = 3;
+
+    // The name of the class that contains the keypoints for this task. This is
+    // used to retrieve the corresponding keypoint indices from the label map.
+    // Note that this corresponds to the "name" field, not "display_name".
+    optional string keypoint_class_name = 4;
+
+    // The standard deviation of the Gaussian kernel used to generate the
+    // keypoint heatmap. The unit is the pixel in the output image. It is to
+    // provide the flexibility of using different sizes of Gaussian kernel for
+    // each keypoint class. Note that if provided, the keypoint standard
+    // deviations will be overridden by the specified values here, otherwise,
+    // the default value 5.0 will be used.
+    // TODO(yuhuic): Update the default value once we found the best value.
+    map<string, float> keypoint_label_to_std = 5;
+
+    // Loss weights corresponding to different heads.
+    optional float keypoint_regression_loss_weight = 6 [default = 1.0];
+    optional float keypoint_heatmap_loss_weight = 7 [default = 1.0];
+    optional float keypoint_offset_loss_weight = 8 [default = 1.0];
+
+    // The initial bias value of the convolution kernel of the keypoint heatmap
+    // prediction head. -2.19 corresponds to predicting foreground with
+    // a probability of 0.1. See "Focal Loss for Dense Object Detection"
+    // at https://arxiv.org/abs/1708.02002.
+    optional float heatmap_bias_init = 9 [default = -2.19];
+
+    // The heatmap score threshold for a keypoint to become a valid candidate.
+    optional float keypoint_candidate_score_threshold = 10 [default = 0.1];
+
+    // The maximum number of candidates to retrieve for each keypoint.
+    optional int32 num_candidates_per_keypoint = 11 [default = 100];
+
+    // Max pool kernel size to use to pull off peak score locations in a
+    // neighborhood (independently for each keypoint types).
+    optional int32 peak_max_pool_kernel_size = 12 [default = 3];
+
+    // The default score to use for regressed keypoints that are not
+    // successfully snapped to a nearby candidate.
+    optional float unmatched_keypoint_score = 13 [default = 0.1];
+
+    // The multiplier to expand the bounding boxes (either the provided boxes or
+    // those which tightly cover the regressed keypoints). Note that new
+    // expanded box for an instance becomes the feasible search window for all
+    // associated keypoints.
+    optional float box_scale = 14 [default = 1.2];
+
+    // The scale parameter that multiplies the largest dimension of a bounding
+    // box. The resulting distance becomes a search radius for candidates in the
+    // vicinity of each regressed keypoint.
+    optional float candidate_search_scale = 15 [default = 0.3];
+
+    // One of ['min_distance', 'score_distance_ratio'] indicating how to select
+    // the keypoint candidate.
+    optional string candidate_ranking_mode = 16 [default = "min_distance"];
+
+    // The radius (in the unit of output pixel) around heatmap peak to assign
+    // the offset targets. If set 0, then the offset target will only be
+    // assigned to the heatmap peak (same behavior as the original paper).
+    optional int32 offset_peak_radius = 17 [default = 0];
+
+    // Indicates whether to assign offsets for each keypoint channel
+    // separately. If set False, the output offset target has the shape
+    // [batch_size, out_height, out_width, 2] (same behavior as the original
+    // paper). If set True, the output offset target has the shape [batch_size,
+    // out_height, out_width, 2 * num_keypoints] (recommended when the
+    // offset_peak_radius is not zero).
+    optional bool per_keypoint_offset = 18 [default = false];
+  }
+  repeated KeypointEstimation keypoint_estimation_task = 7;
+
+  // Parameters which are related to mask estimation task.
+  // Note: Currently, CenterNet supports a weak instance segmentation, where
+  // semantic segmentation masks are estimated, and then cropped based on
+  // bounding box detections. Therefore, it is possible for the same image
+  // pixel to be assigned to multiple instances.
+  message MaskEstimation {
+    // Weight of the task loss. The total loss of the model will be their
+    // summation of task losses weighted by the weights.
+    optional float task_loss_weight = 1 [default = 1.0];
+
+    // Classification loss configuration for segmentation loss.
+    optional ClassificationLoss classification_loss = 2;
+
+    // Each instance mask (one per detection) is cropped and resized (bilinear
+    // resampling) from the predicted segmentation feature map. After
+    // resampling, the masks are binarized with the provided score threshold.
+    optional int32 mask_height = 4 [default = 256];
+    optional int32 mask_width = 5 [default = 256];
+    optional float score_threshold = 6 [default = 0.5];
+
+    // The initial bias value of the convlution kernel of the class heatmap
+    // prediction head. -2.19 corresponds to predicting foreground with
+    // a probability of 0.1.
+    optional float heatmap_bias_init = 3 [default = -2.19];
+  }
+  optional MaskEstimation mask_estimation_task = 8;
+}
+
+message CenterNetFeatureExtractor {
+  optional string type = 1;
+
+  // Channel means to be subtracted from each image channel. If not specified,
+  // we use a default value of 0.
+  repeated float channel_means = 2;
+
+  // Channel standard deviations. Each channel will be normalized by dividing
+  // it by its standard deviation. If not specified, we use a default value
+  // of 1.
+  repeated float channel_stds = 3;
+
+  // If set, will change channel order to be [blue, green, red]. This can be
+  // useful to be compatible with some pre-trained feature extractors.
+  optional bool bgr_ordering = 4 [default = false];
+}

research/object_detection/protos/faster_rcnn.proto

@@ -188,7 +188,7 @@ message Context {
   // Next id: 4
 
   // The maximum number of contextual features per-image, used for padding
-  optional int32 max_num_context_features = 1 [default = 8500];
+  optional int32 max_num_context_features = 1 [default = 2000];
 
   // The bottleneck feature dimension of the attention block.
   optional int32 attention_bottleneck_dimension = 2 [default = 2048];

research/object_detection/protos/model.proto

@@ -2,6 +2,7 @@ syntax = "proto2";
 
 package object_detection.protos;
 
+import "object_detection/protos/center_net.proto";
 import "object_detection/protos/faster_rcnn.proto";
 import "object_detection/protos/ssd.proto";
 
@@ -17,6 +18,7 @@ message DetectionModel {
     // value to a function that builds your model.
     ExperimentalModel experimental_model = 3;
 
+    CenterNet center_net = 4;
   }
 }
 

research/object_detection/samples/configs/ssdlite_mobiledet_gpu_320x320_coco_sync_4x4.config

+# SSDLite with MobileDet-GPU feature extractor.
+# Reference: Xiong & Liu et al., https://arxiv.org/abs/2004.14525
+# Trained on COCO, initialized from scratch.
+#
+# 5.07B MulAdds, 13.11M Parameters.
+# Latencies are 11.0ms (fp32), 3.2ms (fp16) and 2.3ms (int8) on Jetson Xavier,
+# optimized using TensorRT 7.1.
+# Achieves 28.7 mAP on COCO14 minival dataset.
+# Achieves 27.5 mAP on COCO17 val dataset.
+#
+# This config is TPU compatible.
+
+model {
+  ssd {
+    inplace_batchnorm_update: true
+    freeze_batchnorm: false
+    num_classes: 90
+    box_coder {
+      faster_rcnn_box_coder {
+        y_scale: 10.0
+        x_scale: 10.0
+        height_scale: 5.0
+        width_scale: 5.0
+      }
+    }
+    matcher {
+      argmax_matcher {
+        matched_threshold: 0.5
+        unmatched_threshold: 0.5
+        ignore_thresholds: false
+        negatives_lower_than_unmatched: true
+        force_match_for_each_row: true
+        use_matmul_gather: true
+      }
+    }
+    similarity_calculator {
+      iou_similarity {
+      }
+    }
+    encode_background_as_zeros: true
+    anchor_generator {
+      ssd_anchor_generator {
+        num_layers: 6
+        min_scale: 0.2
+        max_scale: 0.95
+        aspect_ratios: 1.0
+        aspect_ratios: 2.0
+        aspect_ratios: 0.5
+        aspect_ratios: 3.0
+        aspect_ratios: 0.3333
+      }
+    }
+    image_resizer {
+      fixed_shape_resizer {
+        height: 320
+        width: 320
+      }
+    }
+    box_predictor {
+      convolutional_box_predictor {
+        min_depth: 0
+        max_depth: 0
+        num_layers_before_predictor: 0
+        use_dropout: false
+        dropout_keep_probability: 0.8
+        kernel_size: 3
+        use_depthwise: true
+        box_code_size: 4
+        apply_sigmoid_to_scores: false
+        class_prediction_bias_init: -4.6
+        conv_hyperparams {
+          activation: RELU_6,
+          regularizer {
+            l2_regularizer {
+              weight: 0.00004
+            }
+          }
+          initializer {
+            random_normal_initializer {
+              stddev: 0.03
+              mean: 0.0
+            }
+          }
+          batch_norm {
+            train: true,
+            scale: true,
+            center: true,
+            decay: 0.97,
+            epsilon: 0.001,
+          }
+        }
+      }
+    }
+    feature_extractor {
+      type: 'ssd_mobiledet_gpu'
+      min_depth: 16
+      depth_multiplier: 1.0
+      use_depthwise: true
+      conv_hyperparams {
+        activation: RELU_6,
+        regularizer {
+          l2_regularizer {
+            weight: 0.00004
+          }
+        }
+        initializer {
+          truncated_normal_initializer {
+            stddev: 0.03
+            mean: 0.0
+          }
+        }
+        batch_norm {
+          train: true,
+          scale: true,
+          center: true,
+          decay: 0.97,
+          epsilon: 0.001,
+        }
+      }
+      override_base_feature_extractor_hyperparams: false
+    }
+    loss {
+      classification_loss {
+        weighted_sigmoid_focal {
+          alpha: 0.75,
+          gamma: 2.0
+        }
+      }
+      localization_loss {
+        weighted_smooth_l1 {
+          delta: 1.0
+        }
+      }
+      classification_weight: 1.0
+      localization_weight: 1.0
+    }
+    normalize_loss_by_num_matches: true
+    normalize_loc_loss_by_codesize: true
+    post_processing {
+      batch_non_max_suppression {
+        score_threshold: 1e-8
+        iou_threshold: 0.6
+        max_detections_per_class: 100
+        max_total_detections: 100
+        use_static_shapes: true
+      }
+      score_converter: SIGMOID
+    }
+  }
+}
+
+train_config: {
+  batch_size: 512
+  sync_replicas: true
+  startup_delay_steps: 0
+  replicas_to_aggregate: 32
+  num_steps: 400000
+  data_augmentation_options {
+    random_horizontal_flip {
+    }
+  }
+  data_augmentation_options {
+    ssd_random_crop {
+    }
+  }
+  optimizer {
+    momentum_optimizer: {
+      learning_rate: {
+        cosine_decay_learning_rate {
+          learning_rate_base: 0.8
+          total_steps: 400000
+          warmup_learning_rate: 0.13333
+          warmup_steps: 2000
+        }
+      }
+      momentum_optimizer_value: 0.9
+    }
+    use_moving_average: false
+  }
+  max_number_of_boxes: 100
+  unpad_groundtruth_tensors: false
+}
+
+train_input_reader: {
+  label_map_path: "PATH_TO_BE_CONFIGURED/mscoco_label_map.pbtxt"
+  tf_record_input_reader {
+    input_path: "PATH_TO_BE_CONFIGURED/mscoco_train.record-?????-of-00100"
+  }
+}
+
+eval_config: {
+  metrics_set: "coco_detection_metrics"
+  use_moving_averages: false
+  num_examples: 8000
+}
+
+eval_input_reader: {
+  label_map_path: "PATH_TO_BE_CONFIGURED/mscoco_label_map.pbtxt"
+  shuffle: false
+  num_epochs: 1
+  tf_record_input_reader {
+    input_path: "PATH_TO_BE_CONFIGURED/mscoco_val.record-?????-of-00010"
+  }
+}

research/object_detection/test_images/snapshot_serengeti/README.md

+# Citation and license
+
+The images and metadata in this folder come from the Snapshot Serengeti dataset,
+and were accessed via [LILA.science](http://lila.science/datasets/snapshot-serengeti).
+The images and species-level labels are described in more detail in the
+associated manuscript:
+
+```
+Swanson AB, Kosmala M, Lintott CJ, Simpson RJ, Smith A, Packer C (2015)
+Snapshot Serengeti, high-frequency annotated camera trap images of 40 mammalian
+species in an African savanna. Scientific Data 2: 150026. (DOI) (bibtex)
+```
+
+Please cite this manuscript if you use this dataset.
+
+This data set is released under the
+[Community Data License Agreement (permissive variant)](https://cdla.io/permissive-1-0/).