Refactor tests for Object Detection API. (#8688)

Internal changes

--

PiperOrigin-RevId: 316837667

research/object_detection/models/ssd_feature_extractor_test.py

@@ -31,6 +31,7 @@ from google.protobuf import text_format
 from object_detection.builders import hyperparams_builder
 from object_detection.protos import hyperparams_pb2
 from object_detection.utils import test_case
+from object_detection.utils import test_utils
 
 
 class SsdFeatureExtractorTestBase(test_case.TestCase):
@@ -89,14 +90,13 @@ class SsdFeatureExtractorTestBase(test_case.TestCase):
     """
     pass
 
-  def _extract_features(self,
-                        image_tensor,
-                        depth_multiplier,
-                        pad_to_multiple,
-                        use_explicit_padding=False,
-                        use_depthwise=False,
-                        num_layers=6,
-                        use_keras=False):
+  def _create_features(self,
+                       depth_multiplier,
+                       pad_to_multiple,
+                       use_explicit_padding=False,
+                       use_depthwise=False,
+                       num_layers=6,
+                       use_keras=False):
     kwargs = {}
     if use_explicit_padding:
       kwargs.update({'use_explicit_padding': use_explicit_padding})
@@ -110,6 +110,12 @@ class SsdFeatureExtractorTestBase(test_case.TestCase):
         depth_multiplier,
         pad_to_multiple,
         **kwargs)
+    return feature_extractor
+
+  def _extract_features(self,
+                        image_tensor,
+                        feature_extractor,
+                        use_keras=False):
     if use_keras:
       feature_maps = feature_extractor(image_tensor)
     else:
@@ -127,10 +133,8 @@ class SsdFeatureExtractorTestBase(test_case.TestCase):
                                                    num_layers=6,
                                                    use_keras=False,
                                                    use_depthwise=False):
-
-    def graph_fn(image_tensor):
-      return self._extract_features(
-          image_tensor,
+    with test_utils.GraphContextOrNone() as g:
+      feature_extractor = self._create_features(
           depth_multiplier,
           pad_to_multiple,
           use_explicit_padding=use_explicit_padding,
@@ -138,9 +142,15 @@ class SsdFeatureExtractorTestBase(test_case.TestCase):
           use_keras=use_keras,
           use_depthwise=use_depthwise)
 
+    def graph_fn(image_tensor):
+      return self._extract_features(
+          image_tensor,
+          feature_extractor,
+          use_keras=use_keras)
+
     image_tensor = np.random.rand(batch_size, image_height, image_width,
                                   3).astype(np.float32)
-    feature_maps = self.execute(graph_fn, [image_tensor])
+    feature_maps = self.execute(graph_fn, [image_tensor], graph=g)
     for feature_map, expected_shape in zip(
         feature_maps, expected_feature_map_shapes):
       self.assertAllEqual(feature_map.shape, expected_shape)
@@ -158,11 +168,8 @@ class SsdFeatureExtractorTestBase(test_case.TestCase):
       use_keras=False,
       use_depthwise=False):
 
-    def graph_fn(image_height, image_width):
-      image_tensor = tf.random_uniform([batch_size, image_height, image_width,
-                                        3], dtype=tf.float32)
-      return self._extract_features(
-          image_tensor,
+    with test_utils.GraphContextOrNone() as g:
+      feature_extractor = self._create_features(
           depth_multiplier,
           pad_to_multiple,
           use_explicit_padding=use_explicit_padding,
@@ -170,10 +177,18 @@ class SsdFeatureExtractorTestBase(test_case.TestCase):
           use_keras=use_keras,
           use_depthwise=use_depthwise)
 
+    def graph_fn(image_height, image_width):
+      image_tensor = tf.random_uniform([batch_size, image_height, image_width,
+                                        3], dtype=tf.float32)
+      return self._extract_features(
+          image_tensor,
+          feature_extractor,
+          use_keras=use_keras)
+
     feature_maps = self.execute_cpu(graph_fn, [
         np.array(image_height, dtype=np.int32),
         np.array(image_width, dtype=np.int32)
-    ])
+    ], graph=g)
     for feature_map, expected_shape in zip(
         feature_maps, expected_feature_map_shapes):
       self.assertAllEqual(feature_map.shape, expected_shape)
@@ -186,19 +201,33 @@ class SsdFeatureExtractorTestBase(test_case.TestCase):
       pad_to_multiple,
       use_keras=False,
       use_depthwise=False):
-    preprocessed_inputs = tf.placeholder(tf.float32, (4, None, None, 3))
-    feature_maps = self._extract_features(
-        preprocessed_inputs,
-        depth_multiplier,
-        pad_to_multiple,
-        use_keras=use_keras,
-        use_depthwise=use_depthwise)
-    test_preprocessed_image = np.random.rand(4, image_height, image_width, 3)
-    with self.test_session() as sess:
-      sess.run(tf.global_variables_initializer())
+
+    with test_utils.GraphContextOrNone() as g:
+      batch = 4
+      width = tf.random.uniform([], minval=image_width, maxval=image_width+1,
+                                dtype=tf.int32)
+      height = tf.random.uniform([], minval=image_height, maxval=image_height+1,
+                                 dtype=tf.int32)
+      shape = tf.stack([batch, height, width, 3])
+      preprocessed_inputs = tf.random.uniform(shape)
+      feature_extractor = self._create_features(
+          depth_multiplier,
+          pad_to_multiple,
+          use_keras=use_keras,
+          use_depthwise=use_depthwise)
+
+    def graph_fn():
+      feature_maps = self._extract_features(
+          preprocessed_inputs,
+          feature_extractor,
+          use_keras=use_keras)
+      return feature_maps
+    if self.is_tf2():
+      with self.assertRaises(ValueError):
+        self.execute_cpu(graph_fn, [], graph=g)
+    else:
       with self.assertRaises(tf.errors.InvalidArgumentError):
-        sess.run(feature_maps,
-                 feed_dict={preprocessed_inputs: test_preprocessed_image})
+        self.execute_cpu(graph_fn, [], graph=g)
 
   def check_feature_extractor_variables_under_scope(self,
                                                     depth_multiplier,
@@ -221,11 +250,14 @@ class SsdFeatureExtractorTestBase(test_case.TestCase):
                                       use_depthwise=False):
     g = tf.Graph()
     with g.as_default():
-      preprocessed_inputs = tf.placeholder(tf.float32, (4, None, None, 3))
-      self._extract_features(
-          preprocessed_inputs,
+      feature_extractor = self._create_features(
           depth_multiplier,
           pad_to_multiple,
           use_keras=use_keras,
           use_depthwise=use_depthwise)
+      preprocessed_inputs = tf.placeholder(tf.float32, (4, None, None, 3))
+      self._extract_features(
+          preprocessed_inputs,
+          feature_extractor,
+          use_keras=use_keras)
       return g.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)

research/object_detection/models/ssd_inception_v2_feature_extractor_test.py → research/object_detection/models/ssd_inception_v2_feature_extractor_tf1_test.py

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

research/object_detection/models/ssd_inception_v3_feature_extractor_test.py → research/object_detection/models/ssd_inception_v3_feature_extractor_tf1_test.py

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

research/object_detection/models/ssd_mobiledet_feature_extractor.py

@@ -290,6 +290,72 @@ def mobiledet_edgetpu_backbone(h, multiplier=1.0):
   return endpoints
 
 
+def mobiledet_gpu_backbone(h, multiplier=1.0):
+  """Build a MobileDet GPU backbone."""
+
+  def _scale(filters):
+    return _scale_filters(filters, multiplier)
+
+  ibn = functools.partial(_inverted_bottleneck, activation_fn=tf.nn.relu6)
+  fused = functools.partial(_fused_conv, activation_fn=tf.nn.relu6)
+  tucker = functools.partial(_tucker_conv, activation_fn=tf.nn.relu6)
+
+  endpoints = {}
+  # block 0
+  h = _conv(h, _scale(32), 3, strides=2, activation_fn=tf.nn.relu6)
+
+  # block 1
+  h = tucker(
+      h,
+      _scale(16),
+      input_rank_ratio=0.25,
+      output_rank_ratio=0.25,
+      residual=False)
+  endpoints['C1'] = h
+
+  # block 2
+  h = fused(h, _scale(32), expansion=8, strides=2, residual=False)
+  h = tucker(h, _scale(32), input_rank_ratio=0.25, output_rank_ratio=0.25)
+  h = tucker(h, _scale(32), input_rank_ratio=0.25, output_rank_ratio=0.25)
+  h = tucker(h, _scale(32), input_rank_ratio=0.25, output_rank_ratio=0.25)
+  endpoints['C2'] = h
+
+  # block 3
+  h = fused(
+      h, _scale(64), expansion=8, kernel_size=3, strides=2, residual=False)
+  h = fused(h, _scale(64), expansion=8)
+  h = fused(h, _scale(64), expansion=8)
+  h = fused(h, _scale(64), expansion=4)
+  endpoints['C3'] = h
+
+  # block 4
+  h = fused(
+      h, _scale(128), expansion=8, kernel_size=3, strides=2, residual=False)
+  h = fused(h, _scale(128), expansion=4)
+  h = fused(h, _scale(128), expansion=4)
+  h = fused(h, _scale(128), expansion=4)
+
+  # block 5
+  h = fused(
+      h, _scale(128), expansion=8, kernel_size=3, strides=1, residual=False)
+  h = fused(h, _scale(128), expansion=8)
+  h = fused(h, _scale(128), expansion=8)
+  h = fused(h, _scale(128), expansion=8)
+  endpoints['C4'] = h
+
+  # block 6
+  h = fused(
+      h, _scale(128), expansion=4, kernel_size=3, strides=2, residual=False)
+  h = fused(h, _scale(128), expansion=4)
+  h = fused(h, _scale(128), expansion=4)
+  h = fused(h, _scale(128), expansion=4)
+
+  # block 7
+  h = ibn(h, _scale(384), expansion=8, kernel_size=3, strides=1, residual=False)
+  endpoints['C5'] = h
+  return endpoints
+
+
 class SSDMobileDetFeatureExtractorBase(ssd_meta_arch.SSDFeatureExtractor):
   """Base class of SSD feature extractor using MobileDet features."""
 
@@ -490,3 +556,31 @@ class SSDMobileDetEdgeTPUFeatureExtractor(SSDMobileDetFeatureExtractorBase):
         use_depthwise=use_depthwise,
         override_base_feature_extractor_hyperparams=override_base_feature_extractor_hyperparams,
         scope_name=scope_name)
+
+
+class SSDMobileDetGPUFeatureExtractor(SSDMobileDetFeatureExtractorBase):
+  """MobileDet-GPU feature extractor."""
+
+  def __init__(self,
+               is_training,
+               depth_multiplier,
+               min_depth,
+               pad_to_multiple,
+               conv_hyperparams_fn,
+               reuse_weights=None,
+               use_explicit_padding=False,
+               use_depthwise=False,
+               override_base_feature_extractor_hyperparams=False,
+               scope_name='MobileDetGPU'):
+    super(SSDMobileDetGPUFeatureExtractor, self).__init__(
+        backbone_fn=mobiledet_gpu_backbone,
+        is_training=is_training,
+        depth_multiplier=depth_multiplier,
+        min_depth=min_depth,
+        pad_to_multiple=pad_to_multiple,
+        conv_hyperparams_fn=conv_hyperparams_fn,
+        reuse_weights=reuse_weights,
+        use_explicit_padding=use_explicit_padding,
+        use_depthwise=use_depthwise,
+        override_base_feature_extractor_hyperparams=override_base_feature_extractor_hyperparams,
+        scope_name=scope_name)

research/object_detection/models/ssd_mobiledet_feature_extractor_test.py → research/object_detection/models/ssd_mobiledet_feature_extractor_tf1_test.py

@@ -13,14 +13,20 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for ssd_mobiledet_feature_extractor."""
-
+import unittest
 import tensorflow.compat.v1 as tf
 
-from tensorflow.contrib import quantize as contrib_quantize
 from object_detection.models import ssd_feature_extractor_test
 from object_detection.models import ssd_mobiledet_feature_extractor
+from object_detection.utils import tf_version
+
+try:
+  from tensorflow.contrib import quantize as contrib_quantize  # pylint: disable=g-import-not-at-top
+except:  # pylint: disable=bare-except
+  pass
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SSDMobileDetFeatureExtractorTest(
     ssd_feature_extractor_test.SsdFeatureExtractorTestBase):
 
@@ -105,6 +111,19 @@ class SSDMobileDetFeatureExtractorTest(
     for expected_shape, x in zip(expected_feature_map_shapes, feature_maps):
       self.assertTrue(x.shape.is_compatible_with(expected_shape))
 
+  def test_mobiledet_gpu_returns_correct_shapes(self):
+    expected_feature_map_shapes = [(2, 40, 20, 128), (2, 20, 10, 384),
+                                   (2, 10, 5, 512), (2, 5, 3, 256),
+                                   (2, 3, 2, 256), (2, 2, 1, 128)]
+    feature_extractor = self._create_feature_extractor(
+        ssd_mobiledet_feature_extractor.SSDMobileDetGPUFeatureExtractor)
+    image = tf.random.normal((2, 640, 320, 3))
+    feature_maps = feature_extractor.extract_features(image)
+
+    self.assertEqual(len(expected_feature_map_shapes), len(feature_maps))
+    for expected_shape, x in zip(expected_feature_map_shapes, feature_maps):
+      self.assertTrue(x.shape.is_compatible_with(expected_shape))
+
   def _check_quantization(self, model_fn):
     checkpoint_dir = self.get_temp_dir()
 

research/object_detection/models/ssd_mobilenet_edgetpu_feature_extractor_test.py → research/object_detection/models/ssd_mobilenet_edgetpu_feature_extractor_tf1_test.py

@@ -13,13 +13,15 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for ssd_mobilenet_edgetpu_feature_extractor."""
-
+import unittest
 import tensorflow.compat.v1 as tf
 
 from object_detection.models import ssd_mobilenet_edgetpu_feature_extractor
 from object_detection.models import ssd_mobilenet_edgetpu_feature_extractor_testbase
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SsdMobilenetEdgeTPUFeatureExtractorTest(
     ssd_mobilenet_edgetpu_feature_extractor_testbase
     ._SsdMobilenetEdgeTPUFeatureExtractorTestBase):

research/object_detection/models/ssd_mobilenet_v1_feature_extractor_test.py → research/object_detection/models/ssd_mobilenet_v1_feature_extractor_tf1_test.py

@@ -17,20 +17,16 @@
 By using parameterized test decorator, this test serves for both Slim-based and
 Keras-based Mobilenet V1 feature extractors in SSD.
 """
-from absl.testing import parameterized
-
+import unittest
 import numpy as np
 import tensorflow.compat.v1 as tf
 
 from object_detection.models import ssd_feature_extractor_test
 from object_detection.models import ssd_mobilenet_v1_feature_extractor
-from object_detection.models import ssd_mobilenet_v1_keras_feature_extractor
+from object_detection.utils import tf_version
 
 
-@parameterized.parameters(
-    {'use_keras': False},
-    {'use_keras': True},
-)
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SsdMobilenetV1FeatureExtractorTest(
     ssd_feature_extractor_test.SsdFeatureExtractorTestBase):
 
@@ -59,31 +55,17 @@ class SsdMobilenetV1FeatureExtractorTest(
       an ssd_meta_arch.SSDFeatureExtractor object.
     """
     min_depth = 32
-    if use_keras:
-      return (ssd_mobilenet_v1_keras_feature_extractor
-              .SSDMobileNetV1KerasFeatureExtractor(
-                  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,
-                  use_explicit_padding=use_explicit_padding,
-                  num_layers=num_layers,
-                  name='MobilenetV1'))
-    else:
-      return ssd_mobilenet_v1_feature_extractor.SSDMobileNetV1FeatureExtractor(
-          is_training,
-          depth_multiplier,
-          min_depth,
-          pad_to_multiple,
-          self.conv_hyperparams_fn,
-          use_explicit_padding=use_explicit_padding,
-          num_layers=num_layers)
+    del use_keras
+    return ssd_mobilenet_v1_feature_extractor.SSDMobileNetV1FeatureExtractor(
+        is_training,
+        depth_multiplier,
+        min_depth,
+        pad_to_multiple,
+        self.conv_hyperparams_fn,
+        use_explicit_padding=use_explicit_padding,
+        num_layers=num_layers)
 
-  def test_extract_features_returns_correct_shapes_128(self, use_keras):
+  def test_extract_features_returns_correct_shapes_128(self):
     image_height = 128
     image_width = 128
     depth_multiplier = 1.0
@@ -99,7 +81,7 @@ class SsdMobilenetV1FeatureExtractorTest(
         pad_to_multiple,
         expected_feature_map_shape,
         use_explicit_padding=False,
-        use_keras=use_keras)
+        use_keras=False)
     self.check_extract_features_returns_correct_shape(
         2,
         image_height,
@@ -108,9 +90,9 @@ class SsdMobilenetV1FeatureExtractorTest(
         pad_to_multiple,
         expected_feature_map_shape,
         use_explicit_padding=True,
-        use_keras=use_keras)
+        use_keras=False)
 
-  def test_extract_features_returns_correct_shapes_299(self, use_keras):
+  def test_extract_features_returns_correct_shapes_299(self):
     image_height = 299
     image_width = 299
     depth_multiplier = 1.0
@@ -126,7 +108,7 @@ class SsdMobilenetV1FeatureExtractorTest(
         pad_to_multiple,
         expected_feature_map_shape,
         use_explicit_padding=False,
-        use_keras=use_keras)
+        use_keras=False)
     self.check_extract_features_returns_correct_shape(
         2,
         image_height,
@@ -135,9 +117,9 @@ class SsdMobilenetV1FeatureExtractorTest(
         pad_to_multiple,
         expected_feature_map_shape,
         use_explicit_padding=True,
-        use_keras=use_keras)
+        use_keras=False)
 
-  def test_extract_features_with_dynamic_image_shape(self, use_keras):
+  def test_extract_features_with_dynamic_image_shape(self):
     image_height = 128
     image_width = 128
     depth_multiplier = 1.0
@@ -153,7 +135,7 @@ class SsdMobilenetV1FeatureExtractorTest(
         pad_to_multiple,
         expected_feature_map_shape,
         use_explicit_padding=False,
-        use_keras=use_keras)
+        use_keras=False)
     self.check_extract_features_returns_correct_shape(
         2,
         image_height,
@@ -162,10 +144,10 @@ class SsdMobilenetV1FeatureExtractorTest(
         pad_to_multiple,
         expected_feature_map_shape,
         use_explicit_padding=True,
-        use_keras=use_keras)
+        use_keras=False)
 
   def test_extract_features_returns_correct_shapes_enforcing_min_depth(
-      self, use_keras):
+      self):
     image_height = 299
     image_width = 299
     depth_multiplier = 0.5**12
@@ -181,7 +163,7 @@ class SsdMobilenetV1FeatureExtractorTest(
         pad_to_multiple,
         expected_feature_map_shape,
         use_explicit_padding=False,
-        use_keras=use_keras)
+        use_keras=False)
     self.check_extract_features_returns_correct_shape(
         2,
         image_height,
@@ -190,10 +172,10 @@ class SsdMobilenetV1FeatureExtractorTest(
         pad_to_multiple,
         expected_feature_map_shape,
         use_explicit_padding=True,
-        use_keras=use_keras)
+        use_keras=False)
 
   def test_extract_features_returns_correct_shapes_with_pad_to_multiple(
-      self, use_keras):
+      self):
     image_height = 299
     image_width = 299
     depth_multiplier = 1.0
@@ -209,7 +191,7 @@ class SsdMobilenetV1FeatureExtractorTest(
         pad_to_multiple,
         expected_feature_map_shape,
         use_explicit_padding=False,
-        use_keras=use_keras)
+        use_keras=False)
     self.check_extract_features_returns_correct_shape(
         2,
         image_height,
@@ -218,10 +200,10 @@ class SsdMobilenetV1FeatureExtractorTest(
         pad_to_multiple,
         expected_feature_map_shape,
         use_explicit_padding=True,
-        use_keras=use_keras)
+        use_keras=False)
 
   def test_extract_features_raises_error_with_invalid_image_size(
-      self, use_keras):
+      self):
     image_height = 32
     image_width = 32
     depth_multiplier = 1.0
@@ -231,34 +213,34 @@ class SsdMobilenetV1FeatureExtractorTest(
         image_width,
         depth_multiplier,
         pad_to_multiple,
-        use_keras=use_keras)
+        use_keras=False)
 
-  def test_preprocess_returns_correct_value_range(self, use_keras):
+  def test_preprocess_returns_correct_value_range(self):
     image_height = 128
     image_width = 128
     depth_multiplier = 1
     pad_to_multiple = 1
     test_image = np.random.rand(2, image_height, image_width, 3)
     feature_extractor = self._create_feature_extractor(
-        depth_multiplier, pad_to_multiple, use_keras=use_keras)
+        depth_multiplier, pad_to_multiple, use_keras=False)
     preprocessed_image = feature_extractor.preprocess(test_image)
     self.assertTrue(np.all(np.less_equal(np.abs(preprocessed_image), 1.0)))
 
-  def test_variables_only_created_in_scope(self, use_keras):
+  def test_variables_only_created_in_scope(self):
     depth_multiplier = 1
     pad_to_multiple = 1
     scope_name = 'MobilenetV1'
     self.check_feature_extractor_variables_under_scope(
-        depth_multiplier, pad_to_multiple, scope_name, use_keras=use_keras)
+        depth_multiplier, pad_to_multiple, scope_name, use_keras=False)
 
-  def test_variable_count(self, use_keras):
+  def test_variable_count(self):
     depth_multiplier = 1
     pad_to_multiple = 1
     variables = self.get_feature_extractor_variables(
-        depth_multiplier, pad_to_multiple, use_keras=use_keras)
+        depth_multiplier, pad_to_multiple, use_keras=False)
     self.assertEqual(len(variables), 151)
 
-  def test_has_fused_batchnorm(self, use_keras):
+  def test_has_fused_batchnorm(self):
     image_height = 40
     image_width = 40
     depth_multiplier = 1
@@ -266,17 +248,14 @@ class SsdMobilenetV1FeatureExtractorTest(
     image_placeholder = tf.placeholder(tf.float32,
                                        [1, image_height, image_width, 3])
     feature_extractor = self._create_feature_extractor(
-        depth_multiplier, pad_to_multiple, use_keras=use_keras)
+        depth_multiplier, pad_to_multiple, use_keras=False)
     preprocessed_image = feature_extractor.preprocess(image_placeholder)
-    if use_keras:
-      _ = feature_extractor(preprocessed_image)
-    else:
-      _ = feature_extractor.extract_features(preprocessed_image)
+    _ = feature_extractor.extract_features(preprocessed_image)
     self.assertTrue(
         any('FusedBatchNorm' in op.type
             for op in tf.get_default_graph().get_operations()))
 
-  def test_extract_features_with_fewer_layers(self, use_keras):
+  def test_extract_features_with_fewer_layers(self):
     image_height = 128
     image_width = 128
     depth_multiplier = 1.0
@@ -286,7 +265,7 @@ class SsdMobilenetV1FeatureExtractorTest(
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=False, num_layers=4,
-        use_keras=use_keras)
+        use_keras=False)
 
 
 if __name__ == '__main__':

research/object_detection/models/ssd_mobilenet_v1_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 Mobilenet V1 feature extractors.
+
+By using parameterized test decorator, this test serves for both Slim-based and
+Keras-based Mobilenet V1 feature extractors in SSD.
+"""
+import unittest
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+from object_detection.models import ssd_feature_extractor_test
+from object_detection.models import ssd_mobilenet_v1_keras_feature_extractor
+from object_detection.utils import tf_version
+
+
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
+class SsdMobilenetV1FeatureExtractorTest(
+    ssd_feature_extractor_test.SsdFeatureExtractorTestBase):
+
+  def _create_feature_extractor(self,
+                                depth_multiplier,
+                                pad_to_multiple,
+                                use_explicit_padding=False,
+                                num_layers=6,
+                                is_training=False,
+                                use_keras=False):
+    """Constructs a new feature extractor.
+
+    Args:
+      depth_multiplier: float depth multiplier for feature extractor
+      pad_to_multiple: the nearest multiple to zero pad the input height and
+        width dimensions to.
+      use_explicit_padding: Use 'VALID' padding for convolutions, but prepad
+        inputs so that the output dimensions are the same as if 'SAME' padding
+        were used.
+      num_layers: number of SSD layers.
+      is_training: whether the network is in training mode.
+      use_keras: if True builds a keras-based feature extractor, if False builds
+        a slim-based one.
+
+    Returns:
+      an ssd_meta_arch.SSDFeatureExtractor object.
+    """
+    del use_keras
+    min_depth = 32
+    return (ssd_mobilenet_v1_keras_feature_extractor
+            .SSDMobileNetV1KerasFeatureExtractor(
+                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,
+                use_explicit_padding=use_explicit_padding,
+                num_layers=num_layers,
+                name='MobilenetV1'))
+
+  def test_extract_features_returns_correct_shapes_128(self):
+    image_height = 128
+    image_width = 128
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 8, 8, 512), (2, 4, 4, 1024),
+                                  (2, 2, 2, 512), (2, 1, 1, 256),
+                                  (2, 1, 1, 256), (2, 1, 1, 128)]
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=False,
+        use_keras=True)
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=True,
+        use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_299(self):
+    image_height = 299
+    image_width = 299
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 19, 19, 512), (2, 10, 10, 1024),
+                                  (2, 5, 5, 512), (2, 3, 3, 256),
+                                  (2, 2, 2, 256), (2, 1, 1, 128)]
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=False,
+        use_keras=True)
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=True,
+        use_keras=True)
+
+  def test_extract_features_with_dynamic_image_shape(self):
+    image_height = 128
+    image_width = 128
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 8, 8, 512), (2, 4, 4, 1024),
+                                  (2, 2, 2, 512), (2, 1, 1, 256),
+                                  (2, 1, 1, 256), (2, 1, 1, 128)]
+    self.check_extract_features_returns_correct_shapes_with_dynamic_inputs(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=False,
+        use_keras=True)
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=True,
+        use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_enforcing_min_depth(
+      self):
+    image_height = 299
+    image_width = 299
+    depth_multiplier = 0.5**12
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 19, 19, 32), (2, 10, 10, 32),
+                                  (2, 5, 5, 32), (2, 3, 3, 32), (2, 2, 2, 32),
+                                  (2, 1, 1, 32)]
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=False,
+        use_keras=True)
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=True,
+        use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_with_pad_to_multiple(
+      self):
+    image_height = 299
+    image_width = 299
+    depth_multiplier = 1.0
+    pad_to_multiple = 32
+    expected_feature_map_shape = [(2, 20, 20, 512), (2, 10, 10, 1024),
+                                  (2, 5, 5, 512), (2, 3, 3, 256),
+                                  (2, 2, 2, 256), (2, 1, 1, 128)]
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=False,
+        use_keras=True)
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=True,
+        use_keras=True)
+
+  def test_extract_features_raises_error_with_invalid_image_size(
+      self):
+    image_height = 32
+    image_width = 32
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    self.check_extract_features_raises_error_with_invalid_image_size(
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        use_keras=True)
+
+  def test_preprocess_returns_correct_value_range(self):
+    image_height = 128
+    image_width = 128
+    depth_multiplier = 1
+    pad_to_multiple = 1
+    test_image = np.random.rand(2, image_height, image_width, 3)
+    feature_extractor = self._create_feature_extractor(
+        depth_multiplier, pad_to_multiple, use_keras=True)
+    preprocessed_image = feature_extractor.preprocess(test_image)
+    self.assertTrue(np.all(np.less_equal(np.abs(preprocessed_image), 1.0)))
+
+  def test_extract_features_with_fewer_layers(self):
+    image_height = 128
+    image_width = 128
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 8, 8, 512), (2, 4, 4, 1024),
+                                  (2, 2, 2, 512), (2, 1, 1, 256)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=False, num_layers=4,
+        use_keras=True)
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/object_detection/models/ssd_mobilenet_v1_fpn_feature_extractor_test.py → research/object_detection/models/ssd_mobilenet_v1_fpn_feature_extractor_tf1_test.py

@@ -18,19 +18,16 @@
 By using parameterized test decorator, this test serves for both Slim-based and
 Keras-based Mobilenet V1 FPN feature extractors in SSD.
 """
-from absl.testing import parameterized
+import unittest
 import numpy as np
 import tensorflow.compat.v1 as tf
 
 from object_detection.models import ssd_feature_extractor_test
 from object_detection.models import ssd_mobilenet_v1_fpn_feature_extractor
-from object_detection.models import ssd_mobilenet_v1_fpn_keras_feature_extractor
+from object_detection.utils import tf_version
 
 
-@parameterized.parameters(
-    {'use_keras': False},
-    {'use_keras': True},
-)
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SsdMobilenetV1FpnFeatureExtractorTest(
     ssd_feature_extractor_test.SsdFeatureExtractorTestBase):
 
@@ -52,33 +49,19 @@ class SsdMobilenetV1FpnFeatureExtractorTest(
     Returns:
       an ssd_meta_arch.SSDFeatureExtractor object.
     """
+    del use_keras
     min_depth = 32
-    if use_keras:
-      return (ssd_mobilenet_v1_fpn_keras_feature_extractor.
-              SSDMobileNetV1FpnKerasFeatureExtractor(
-                  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,
-                  use_explicit_padding=use_explicit_padding,
-                  use_depthwise=True,
-                  name='MobilenetV1_FPN'))
-    else:
-      return (ssd_mobilenet_v1_fpn_feature_extractor.
-              SSDMobileNetV1FpnFeatureExtractor(
-                  is_training,
-                  depth_multiplier,
-                  min_depth,
-                  pad_to_multiple,
-                  self.conv_hyperparams_fn,
-                  use_depthwise=True,
-                  use_explicit_padding=use_explicit_padding))
-
-  def test_extract_features_returns_correct_shapes_256(self, use_keras):
+    return (ssd_mobilenet_v1_fpn_feature_extractor.
+            SSDMobileNetV1FpnFeatureExtractor(
+                is_training,
+                depth_multiplier,
+                min_depth,
+                pad_to_multiple,
+                self.conv_hyperparams_fn,
+                use_depthwise=True,
+                use_explicit_padding=use_explicit_padding))
+
+  def test_extract_features_returns_correct_shapes_256(self):
     image_height = 256
     image_width = 256
     depth_multiplier = 1.0
@@ -89,13 +72,13 @@ class SsdMobilenetV1FpnFeatureExtractorTest(
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=False,
-        use_keras=use_keras)
+        use_keras=False)
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=True,
-        use_keras=use_keras)
+        use_keras=False)
 
-  def test_extract_features_returns_correct_shapes_384(self, use_keras):
+  def test_extract_features_returns_correct_shapes_384(self):
     image_height = 320
     image_width = 320
     depth_multiplier = 1.0
@@ -106,13 +89,13 @@ class SsdMobilenetV1FpnFeatureExtractorTest(
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=False,
-        use_keras=use_keras)
+        use_keras=False)
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=True,
-        use_keras=use_keras)
+        use_keras=False)
 
-  def test_extract_features_with_dynamic_image_shape(self, use_keras):
+  def test_extract_features_with_dynamic_image_shape(self):
     image_height = 256
     image_width = 256
     depth_multiplier = 1.0
@@ -123,14 +106,14 @@ class SsdMobilenetV1FpnFeatureExtractorTest(
     self.check_extract_features_returns_correct_shapes_with_dynamic_inputs(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=False,
-        use_keras=use_keras)
+        use_keras=False)
     self.check_extract_features_returns_correct_shapes_with_dynamic_inputs(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=True,
-        use_keras=use_keras)
+        use_keras=False)
 
   def test_extract_features_returns_correct_shapes_with_pad_to_multiple(
-      self, use_keras):
+      self):
     image_height = 299
     image_width = 299
     depth_multiplier = 1.0
@@ -141,14 +124,14 @@ class SsdMobilenetV1FpnFeatureExtractorTest(
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=False,
-        use_keras=use_keras)
+        use_keras=False)
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=True,
-        use_keras=use_keras)
+        use_keras=False)
 
   def test_extract_features_returns_correct_shapes_enforcing_min_depth(
-      self, use_keras):
+      self):
     image_height = 256
     image_width = 256
     depth_multiplier = 0.5**12
@@ -159,23 +142,23 @@ class SsdMobilenetV1FpnFeatureExtractorTest(
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=False,
-        use_keras=use_keras)
+        use_keras=False)
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
         expected_feature_map_shape, use_explicit_padding=True,
-        use_keras=use_keras)
+        use_keras=False)
 
   def test_extract_features_raises_error_with_invalid_image_size(
-      self, use_keras):
+      self):
     image_height = 32
     image_width = 32
     depth_multiplier = 1.0
     pad_to_multiple = 1
     self.check_extract_features_raises_error_with_invalid_image_size(
         image_height, image_width, depth_multiplier, pad_to_multiple,
-        use_keras=use_keras)
+        use_keras=False)
 
-  def test_preprocess_returns_correct_value_range(self, use_keras):
+  def test_preprocess_returns_correct_value_range(self):
     image_height = 256
     image_width = 256
     depth_multiplier = 1
@@ -183,25 +166,25 @@ class SsdMobilenetV1FpnFeatureExtractorTest(
     test_image = np.random.rand(2, image_height, image_width, 3)
     feature_extractor = self._create_feature_extractor(depth_multiplier,
                                                        pad_to_multiple,
-                                                       use_keras=use_keras)
+                                                       use_keras=False)
     preprocessed_image = feature_extractor.preprocess(test_image)
     self.assertTrue(np.all(np.less_equal(np.abs(preprocessed_image), 1.0)))
 
-  def test_variables_only_created_in_scope(self, use_keras):
+  def test_variables_only_created_in_scope(self):
     depth_multiplier = 1
     pad_to_multiple = 1
     scope_name = 'MobilenetV1'
     self.check_feature_extractor_variables_under_scope(
-        depth_multiplier, pad_to_multiple, scope_name, use_keras=use_keras)
+        depth_multiplier, pad_to_multiple, scope_name, use_keras=False)
 
-  def test_variable_count(self, use_keras):
+  def test_variable_count(self):
     depth_multiplier = 1
     pad_to_multiple = 1
     variables = self.get_feature_extractor_variables(
-        depth_multiplier, pad_to_multiple, use_keras=use_keras)
+        depth_multiplier, pad_to_multiple, use_keras=False)
     self.assertEqual(len(variables), 153)
 
-  def test_fused_batchnorm(self, use_keras):
+  def test_fused_batchnorm(self):
     image_height = 256
     image_width = 256
     depth_multiplier = 1
@@ -210,12 +193,9 @@ class SsdMobilenetV1FpnFeatureExtractorTest(
                                        [1, image_height, image_width, 3])
     feature_extractor = self._create_feature_extractor(depth_multiplier,
                                                        pad_to_multiple,
-                                                       use_keras=use_keras)
+                                                       use_keras=False)
     preprocessed_image = feature_extractor.preprocess(image_placeholder)
-    if use_keras:
-      _ = feature_extractor(preprocessed_image)
-    else:
-      _ = feature_extractor.extract_features(preprocessed_image)
+    _ = feature_extractor.extract_features(preprocessed_image)
 
     self.assertTrue(
         any('FusedBatchNorm' in op.type

research/object_detection/models/ssd_mobilenet_v1_fpn_feature_extractor_tf2_test.py

+# Copyright 2018 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_mobilenet_v1_fpn_feature_extractor.
+
+By using parameterized test decorator, this test serves for both Slim-based and
+Keras-based Mobilenet V1 FPN feature extractors in SSD.
+"""
+import unittest
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+from object_detection.models import ssd_feature_extractor_test
+from object_detection.models import ssd_mobilenet_v1_fpn_keras_feature_extractor
+from object_detection.utils import tf_version
+
+
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
+class SsdMobilenetV1FpnFeatureExtractorTest(
+    ssd_feature_extractor_test.SsdFeatureExtractorTestBase):
+
+  def _create_feature_extractor(self, depth_multiplier, pad_to_multiple,
+                                is_training=True, use_explicit_padding=False,
+                                use_keras=True):
+    """Constructs a new feature extractor.
+
+    Args:
+      depth_multiplier: float depth multiplier for feature extractor
+      pad_to_multiple: the nearest multiple to zero pad the input height and
+        width dimensions to.
+      is_training: whether the network is in training mode.
+      use_explicit_padding: Use 'VALID' padding for convolutions, but prepad
+        inputs so that the output dimensions are the same as if 'SAME' padding
+        were used.
+      use_keras: if True builds a keras-based feature extractor, if False builds
+        a slim-based one.
+    Returns:
+      an ssd_meta_arch.SSDFeatureExtractor object.
+    """
+    min_depth = 32
+    del use_keras
+    return (ssd_mobilenet_v1_fpn_keras_feature_extractor.
+            SSDMobileNetV1FpnKerasFeatureExtractor(
+                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,
+                use_explicit_padding=use_explicit_padding,
+                use_depthwise=True,
+                name='MobilenetV1_FPN'))
+
+  def test_extract_features_returns_correct_shapes_256(self):
+    image_height = 256
+    image_width = 256
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 32, 32, 256), (2, 16, 16, 256),
+                                  (2, 8, 8, 256), (2, 4, 4, 256),
+                                  (2, 2, 2, 256)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=False,
+        use_keras=True)
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=True,
+        use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_384(self):
+    image_height = 320
+    image_width = 320
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 40, 40, 256), (2, 20, 20, 256),
+                                  (2, 10, 10, 256), (2, 5, 5, 256),
+                                  (2, 3, 3, 256)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=False,
+        use_keras=True)
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=True,
+        use_keras=True)
+
+  def test_extract_features_with_dynamic_image_shape(self):
+    image_height = 256
+    image_width = 256
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 32, 32, 256), (2, 16, 16, 256),
+                                  (2, 8, 8, 256), (2, 4, 4, 256),
+                                  (2, 2, 2, 256)]
+    self.check_extract_features_returns_correct_shapes_with_dynamic_inputs(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=False,
+        use_keras=True)
+    self.check_extract_features_returns_correct_shapes_with_dynamic_inputs(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=True,
+        use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_with_pad_to_multiple(
+      self):
+    image_height = 299
+    image_width = 299
+    depth_multiplier = 1.0
+    pad_to_multiple = 32
+    expected_feature_map_shape = [(2, 40, 40, 256), (2, 20, 20, 256),
+                                  (2, 10, 10, 256), (2, 5, 5, 256),
+                                  (2, 3, 3, 256)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=False,
+        use_keras=True)
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=True,
+        use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_enforcing_min_depth(
+      self):
+    image_height = 256
+    image_width = 256
+    depth_multiplier = 0.5**12
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 32, 32, 32), (2, 16, 16, 32),
+                                  (2, 8, 8, 32), (2, 4, 4, 32),
+                                  (2, 2, 2, 32)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=False,
+        use_keras=True)
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=True,
+        use_keras=True)
+
+  def test_extract_features_raises_error_with_invalid_image_size(
+      self):
+    image_height = 32
+    image_width = 32
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    self.check_extract_features_raises_error_with_invalid_image_size(
+        image_height, image_width, depth_multiplier, pad_to_multiple,
+        use_keras=True)
+
+  def test_preprocess_returns_correct_value_range(self):
+    image_height = 256
+    image_width = 256
+    depth_multiplier = 1
+    pad_to_multiple = 1
+    test_image = np.random.rand(2, image_height, image_width, 3)
+    feature_extractor = self._create_feature_extractor(depth_multiplier,
+                                                       pad_to_multiple,
+                                                       use_keras=True)
+    preprocessed_image = feature_extractor.preprocess(test_image)
+    self.assertTrue(np.all(np.less_equal(np.abs(preprocessed_image), 1.0)))
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/object_detection/models/ssd_mobilenet_v1_fpn_keras_feature_extractor.py

@@ -123,7 +123,7 @@ class SSDMobileNetV1FpnKerasFeatureExtractor(
         'Conv2d_3_pointwise', 'Conv2d_5_pointwise', 'Conv2d_11_pointwise',
         'Conv2d_13_pointwise'
     ]
-    self._mobilenet_v1 = None
+    self.classification_backbone = None
     self._fpn_features_generator = None
     self._coarse_feature_layers = []
 
@@ -147,7 +147,7 @@ class SSDMobileNetV1FpnKerasFeatureExtractor(
         name='conv_pw_11_relu').output
     conv2d_13_pointwise = full_mobilenet_v1.get_layer(
         name='conv_pw_13_relu').output
-    self._mobilenet_v1 = tf.keras.Model(
+    self.classification_backbone = tf.keras.Model(
         inputs=full_mobilenet_v1.inputs,
         outputs=[conv2d_3_pointwise, conv2d_5_pointwise,
                  conv2d_11_pointwise, conv2d_13_pointwise]
@@ -218,7 +218,7 @@ class SSDMobileNetV1FpnKerasFeatureExtractor(
     preprocessed_inputs = shape_utils.check_min_image_dim(
         33, preprocessed_inputs)
 
-    image_features = self._mobilenet_v1(
+    image_features = self.classification_backbone(
         ops.pad_to_multiple(preprocessed_inputs, self._pad_to_multiple))
 
     feature_block_list = []
@@ -243,3 +243,14 @@ class SSDMobileNetV1FpnKerasFeatureExtractor(
         last_feature_map = layer(last_feature_map)
       feature_maps.append(last_feature_map)
     return feature_maps
+
+  def restore_from_classification_checkpoint_fn(self, feature_extractor_scope):
+    """Returns a map for restoring from an (object-based) checkpoint.
+
+    Args:
+      feature_extractor_scope: A scope name for the feature extractor (unused).
+
+    Returns:
+      A dict mapping keys to Keras models
+    """
+    return {'feature_extractor': self.classification_backbone}

research/object_detection/models/ssd_mobilenet_v1_keras_feature_extractor.py

@@ -93,7 +93,7 @@ class SSDMobileNetV1KerasFeatureExtractor(
         'use_explicit_padding': self._use_explicit_padding,
         'use_depthwise': self._use_depthwise,
     }
-    self._mobilenet_v1 = None
+    self.classification_backbone = None
     self._feature_map_generator = None
 
   def build(self, input_shape):
@@ -111,7 +111,7 @@ class SSDMobileNetV1KerasFeatureExtractor(
         name='conv_pw_11_relu').output
     conv2d_13_pointwise = full_mobilenet_v1.get_layer(
         name='conv_pw_13_relu').output
-    self._mobilenet_v1 = tf.keras.Model(
+    self.classification_backbone = tf.keras.Model(
         inputs=full_mobilenet_v1.inputs,
         outputs=[conv2d_11_pointwise, conv2d_13_pointwise])
     self._feature_map_generator = (
@@ -155,7 +155,7 @@ class SSDMobileNetV1KerasFeatureExtractor(
     preprocessed_inputs = shape_utils.check_min_image_dim(
         33, preprocessed_inputs)
 
-    image_features = self._mobilenet_v1(
+    image_features = self.classification_backbone(
         ops.pad_to_multiple(preprocessed_inputs, self._pad_to_multiple))
 
     feature_maps = self._feature_map_generator({
@@ -163,3 +163,14 @@ class SSDMobileNetV1KerasFeatureExtractor(
         'Conv2d_13_pointwise': image_features[1]})
 
     return list(feature_maps.values())
+
+  def restore_from_classification_checkpoint_fn(self, feature_extractor_scope):
+    """Returns a map for restoring from an (object-based) checkpoint.
+
+    Args:
+      feature_extractor_scope: A scope name for the feature extractor (unused).
+
+    Returns:
+      A dict mapping keys to Keras models
+    """
+    return {'feature_extractor': self.classification_backbone}

research/object_detection/models/ssd_mobilenet_v1_ppn_feature_extractor_test.py → research/object_detection/models/ssd_mobilenet_v1_ppn_feature_extractor_tf1_test.py

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

research/object_detection/models/ssd_mobilenet_v2_feature_extractor_test.py → research/object_detection/models/ssd_mobilenet_v2_feature_extractor_tf1_test.py

@@ -14,20 +14,17 @@
 # ==============================================================================
 
 """Tests for ssd_mobilenet_v2_feature_extractor."""
-from absl.testing import parameterized
+import unittest
 
 import numpy as np
 import tensorflow.compat.v1 as tf
 
 from object_detection.models import ssd_feature_extractor_test
 from object_detection.models import ssd_mobilenet_v2_feature_extractor
-from object_detection.models import ssd_mobilenet_v2_keras_feature_extractor
+from object_detection.utils import tf_version
 
 
-@parameterized.parameters(
-    {'use_keras': False},
-    {'use_keras': True},
-)
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SsdMobilenetV2FeatureExtractorTest(
     ssd_feature_extractor_test.SsdFeatureExtractorTestBase):
 
@@ -35,8 +32,7 @@ class SsdMobilenetV2FeatureExtractorTest(
                                 depth_multiplier,
                                 pad_to_multiple,
                                 use_explicit_padding=False,
-                                num_layers=6,
-                                use_keras=False):
+                                num_layers=6):
     """Constructs a new feature extractor.
 
     Args:
@@ -47,36 +43,20 @@ class SsdMobilenetV2FeatureExtractorTest(
         inputs so that the output dimensions are the same as if 'SAME' padding
         were used.
       num_layers: number of SSD layers.
-      use_keras: if True builds a keras-based feature extractor, if False builds
-        a slim-based one.
     Returns:
       an ssd_meta_arch.SSDFeatureExtractor object.
     """
     min_depth = 32
-    if use_keras:
-      return (ssd_mobilenet_v2_keras_feature_extractor.
-              SSDMobileNetV2KerasFeatureExtractor(
-                  is_training=False,
-                  depth_multiplier=depth_multiplier,
-                  min_depth=min_depth,
-                  pad_to_multiple=pad_to_multiple,
-                  conv_hyperparams=self._build_conv_hyperparams(),
-                  freeze_batchnorm=False,
-                  inplace_batchnorm_update=False,
-                  use_explicit_padding=use_explicit_padding,
-                  num_layers=num_layers,
-                  name='MobilenetV2'))
-    else:
-      return ssd_mobilenet_v2_feature_extractor.SSDMobileNetV2FeatureExtractor(
-          False,
-          depth_multiplier,
-          min_depth,
-          pad_to_multiple,
-          self.conv_hyperparams_fn,
-          use_explicit_padding=use_explicit_padding,
-          num_layers=num_layers)
-
-  def test_extract_features_returns_correct_shapes_128(self, use_keras):
+    return ssd_mobilenet_v2_feature_extractor.SSDMobileNetV2FeatureExtractor(
+        False,
+        depth_multiplier,
+        min_depth,
+        pad_to_multiple,
+        self.conv_hyperparams_fn,
+        use_explicit_padding=use_explicit_padding,
+        num_layers=num_layers)
+
+  def test_extract_features_returns_correct_shapes_128(self):
     image_height = 128
     image_width = 128
     depth_multiplier = 1.0
@@ -86,10 +66,10 @@ class SsdMobilenetV2FeatureExtractorTest(
                                   (2, 1, 1, 256), (2, 1, 1, 128)]
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape, use_keras=use_keras)
+        expected_feature_map_shape)
 
   def test_extract_features_returns_correct_shapes_128_explicit_padding(
-      self, use_keras):
+      self):
     image_height = 128
     image_width = 128
     depth_multiplier = 1.0
@@ -99,11 +79,10 @@ class SsdMobilenetV2FeatureExtractorTest(
                                   (2, 1, 1, 256), (2, 1, 1, 128)]
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape, use_explicit_padding=True,
-        use_keras=use_keras)
+        expected_feature_map_shape, use_explicit_padding=True)
 
   def test_extract_features_returns_correct_shapes_with_dynamic_inputs(
-      self, use_keras):
+      self):
     image_height = 128
     image_width = 128
     depth_multiplier = 1.0
@@ -113,9 +92,9 @@ class SsdMobilenetV2FeatureExtractorTest(
                                   (2, 1, 1, 256), (2, 1, 1, 128)]
     self.check_extract_features_returns_correct_shapes_with_dynamic_inputs(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape, use_keras=use_keras)
+        expected_feature_map_shape)
 
-  def test_extract_features_returns_correct_shapes_299(self, use_keras):
+  def test_extract_features_returns_correct_shapes_299(self):
     image_height = 299
     image_width = 299
     depth_multiplier = 1.0
@@ -125,10 +104,10 @@ class SsdMobilenetV2FeatureExtractorTest(
                                   (2, 2, 2, 256), (2, 1, 1, 128)]
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape, use_keras=use_keras)
+        expected_feature_map_shape)
 
   def test_extract_features_returns_correct_shapes_enforcing_min_depth(
-      self, use_keras):
+      self):
     image_height = 299
     image_width = 299
     depth_multiplier = 0.5**12
@@ -138,10 +117,10 @@ class SsdMobilenetV2FeatureExtractorTest(
                                   (2, 2, 2, 32), (2, 1, 1, 32)]
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape, use_keras=use_keras)
+        expected_feature_map_shape)
 
   def test_extract_features_returns_correct_shapes_with_pad_to_multiple(
-      self, use_keras):
+      self):
     image_height = 299
     image_width = 299
     depth_multiplier = 1.0
@@ -151,45 +130,43 @@ class SsdMobilenetV2FeatureExtractorTest(
                                   (2, 2, 2, 256), (2, 1, 1, 128)]
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape, use_keras=use_keras)
+        expected_feature_map_shape)
 
   def test_extract_features_raises_error_with_invalid_image_size(
-      self, use_keras):
+      self):
     image_height = 32
     image_width = 32
     depth_multiplier = 1.0
     pad_to_multiple = 1
     self.check_extract_features_raises_error_with_invalid_image_size(
-        image_height, image_width, depth_multiplier, pad_to_multiple,
-        use_keras=use_keras)
+        image_height, image_width, depth_multiplier, pad_to_multiple)
 
-  def test_preprocess_returns_correct_value_range(self, use_keras):
+  def test_preprocess_returns_correct_value_range(self):
     image_height = 128
     image_width = 128
     depth_multiplier = 1
     pad_to_multiple = 1
     test_image = np.random.rand(4, image_height, image_width, 3)
     feature_extractor = self._create_feature_extractor(depth_multiplier,
-                                                       pad_to_multiple,
-                                                       use_keras=use_keras)
+                                                       pad_to_multiple)
     preprocessed_image = feature_extractor.preprocess(test_image)
     self.assertTrue(np.all(np.less_equal(np.abs(preprocessed_image), 1.0)))
 
-  def test_variables_only_created_in_scope(self, use_keras):
+  def test_variables_only_created_in_scope(self):
     depth_multiplier = 1
     pad_to_multiple = 1
     scope_name = 'MobilenetV2'
     self.check_feature_extractor_variables_under_scope(
-        depth_multiplier, pad_to_multiple, scope_name, use_keras=use_keras)
+        depth_multiplier, pad_to_multiple, scope_name)
 
-  def test_variable_count(self, use_keras):
+  def test_variable_count(self):
     depth_multiplier = 1
     pad_to_multiple = 1
     variables = self.get_feature_extractor_variables(
-        depth_multiplier, pad_to_multiple, use_keras=use_keras)
+        depth_multiplier, pad_to_multiple)
     self.assertEqual(len(variables), 292)
 
-  def test_has_fused_batchnorm(self, use_keras):
+  def test_has_fused_batchnorm(self):
     image_height = 40
     image_width = 40
     depth_multiplier = 1
@@ -197,17 +174,13 @@ class SsdMobilenetV2FeatureExtractorTest(
     image_placeholder = tf.placeholder(tf.float32,
                                        [1, image_height, image_width, 3])
     feature_extractor = self._create_feature_extractor(depth_multiplier,
-                                                       pad_to_multiple,
-                                                       use_keras=use_keras)
+                                                       pad_to_multiple)
     preprocessed_image = feature_extractor.preprocess(image_placeholder)
-    if use_keras:
-      _ = feature_extractor(preprocessed_image)
-    else:
-      _ = feature_extractor.extract_features(preprocessed_image)
+    _ = feature_extractor.extract_features(preprocessed_image)
     self.assertTrue(any('FusedBatchNorm' in op.type
                         for op in tf.get_default_graph().get_operations()))
 
-  def test_extract_features_with_fewer_layers(self, use_keras):
+  def test_extract_features_with_fewer_layers(self):
     image_height = 128
     image_width = 128
     depth_multiplier = 1.0
@@ -216,8 +189,7 @@ class SsdMobilenetV2FeatureExtractorTest(
                                   (2, 2, 2, 512), (2, 1, 1, 256)]
     self.check_extract_features_returns_correct_shape(
         2, image_height, image_width, depth_multiplier, pad_to_multiple,
-        expected_feature_map_shape, use_explicit_padding=False, num_layers=4,
-        use_keras=use_keras)
+        expected_feature_map_shape, use_explicit_padding=False, num_layers=4)
 
 
 if __name__ == '__main__':

research/object_detection/models/ssd_mobilenet_v2_feature_extractor_tf2_test.py

+# Copyright 2018 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_mobilenet_v2_feature_extractor."""
+import unittest
+
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+from object_detection.models import ssd_feature_extractor_test
+from object_detection.models import ssd_mobilenet_v2_keras_feature_extractor
+from object_detection.utils import tf_version
+
+
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
+class SsdMobilenetV2FeatureExtractorTest(
+    ssd_feature_extractor_test.SsdFeatureExtractorTestBase):
+
+  def _create_feature_extractor(self,
+                                depth_multiplier,
+                                pad_to_multiple,
+                                use_explicit_padding=False,
+                                num_layers=6,
+                                use_keras=False):
+    """Constructs a new feature extractor.
+
+    Args:
+      depth_multiplier: float depth multiplier for feature extractor
+      pad_to_multiple: the nearest multiple to zero pad the input height and
+        width dimensions to.
+      use_explicit_padding: use 'VALID' padding for convolutions, but prepad
+        inputs so that the output dimensions are the same as if 'SAME' padding
+        were used.
+      num_layers: number of SSD layers.
+      use_keras: unused argument.
+
+    Returns:
+      an ssd_meta_arch.SSDFeatureExtractor object.
+    """
+    del use_keras
+    min_depth = 32
+    return (ssd_mobilenet_v2_keras_feature_extractor.
+            SSDMobileNetV2KerasFeatureExtractor(
+                is_training=False,
+                depth_multiplier=depth_multiplier,
+                min_depth=min_depth,
+                pad_to_multiple=pad_to_multiple,
+                conv_hyperparams=self._build_conv_hyperparams(),
+                freeze_batchnorm=False,
+                inplace_batchnorm_update=False,
+                use_explicit_padding=use_explicit_padding,
+                num_layers=num_layers,
+                name='MobilenetV2'))
+
+  def test_extract_features_returns_correct_shapes_128(self):
+    image_height = 128
+    image_width = 128
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 8, 8, 576), (2, 4, 4, 1280),
+                                  (2, 2, 2, 512), (2, 1, 1, 256),
+                                  (2, 1, 1, 256), (2, 1, 1, 128)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_128_explicit_padding(
+      self):
+    image_height = 128
+    image_width = 128
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 8, 8, 576), (2, 4, 4, 1280),
+                                  (2, 2, 2, 512), (2, 1, 1, 256),
+                                  (2, 1, 1, 256), (2, 1, 1, 128)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=True, use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_with_dynamic_inputs(
+      self):
+    image_height = 128
+    image_width = 128
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 8, 8, 576), (2, 4, 4, 1280),
+                                  (2, 2, 2, 512), (2, 1, 1, 256),
+                                  (2, 1, 1, 256), (2, 1, 1, 128)]
+    self.check_extract_features_returns_correct_shapes_with_dynamic_inputs(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_299(self):
+    image_height = 299
+    image_width = 299
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 19, 19, 576), (2, 10, 10, 1280),
+                                  (2, 5, 5, 512), (2, 3, 3, 256),
+                                  (2, 2, 2, 256), (2, 1, 1, 128)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_enforcing_min_depth(
+      self):
+    image_height = 299
+    image_width = 299
+    depth_multiplier = 0.5**12
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 19, 19, 192), (2, 10, 10, 32),
+                                  (2, 5, 5, 32), (2, 3, 3, 32),
+                                  (2, 2, 2, 32), (2, 1, 1, 32)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_keras=True)
+
+  def test_extract_features_returns_correct_shapes_with_pad_to_multiple(
+      self):
+    image_height = 299
+    image_width = 299
+    depth_multiplier = 1.0
+    pad_to_multiple = 32
+    expected_feature_map_shape = [(2, 20, 20, 576), (2, 10, 10, 1280),
+                                  (2, 5, 5, 512), (2, 3, 3, 256),
+                                  (2, 2, 2, 256), (2, 1, 1, 128)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_keras=True)
+
+  def test_extract_features_raises_error_with_invalid_image_size(
+      self):
+    image_height = 32
+    image_width = 32
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    self.check_extract_features_raises_error_with_invalid_image_size(
+        image_height, image_width, depth_multiplier, pad_to_multiple,
+        use_keras=True)
+
+  def test_preprocess_returns_correct_value_range(self):
+    image_height = 128
+    image_width = 128
+    depth_multiplier = 1
+    pad_to_multiple = 1
+    test_image = np.random.rand(4, image_height, image_width, 3)
+    feature_extractor = self._create_feature_extractor(depth_multiplier,
+                                                       pad_to_multiple)
+    preprocessed_image = feature_extractor.preprocess(test_image)
+    self.assertTrue(np.all(np.less_equal(np.abs(preprocessed_image), 1.0)))
+
+  def test_variables_only_created_in_scope(self):
+    depth_multiplier = 1
+    pad_to_multiple = 1
+    scope_name = 'MobilenetV2'
+    self.check_feature_extractor_variables_under_scope(
+        depth_multiplier, pad_to_multiple, scope_name, use_keras=True)
+
+  def test_variable_count(self):
+    depth_multiplier = 1
+    pad_to_multiple = 1
+    variables = self.get_feature_extractor_variables(
+        depth_multiplier, pad_to_multiple, use_keras=True)
+    self.assertEqual(len(variables), 292)
+
+  def test_extract_features_with_fewer_layers(self):
+    image_height = 128
+    image_width = 128
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 8, 8, 576), (2, 4, 4, 1280),
+                                  (2, 2, 2, 512), (2, 1, 1, 256)]
+    self.check_extract_features_returns_correct_shape(
+        2, image_height, image_width, depth_multiplier, pad_to_multiple,
+        expected_feature_map_shape, use_explicit_padding=False, num_layers=4,
+        use_keras=True)
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/object_detection/models/ssd_mobilenet_v2_fpn_feature_extractor_test.py → research/object_detection/models/ssd_mobilenet_v2_fpn_feature_extractor_tf1_test.py

@@ -18,31 +18,23 @@
 By using parameterized test decorator, this test serves for both Slim-based and
 Keras-based Mobilenet V2 FPN feature extractors in SSD.
 """
+import unittest
 from absl.testing import parameterized
 import numpy as np
 import tensorflow.compat.v1 as tf
 
 from object_detection.models import ssd_feature_extractor_test
 from object_detection.models import ssd_mobilenet_v2_fpn_feature_extractor
-from object_detection.models import ssd_mobilenet_v2_fpn_keras_feature_extractor
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 @parameterized.parameters(
     {
-        'use_depthwise': False,
-        'use_keras': True
+        'use_depthwise': False
     },
     {
-        'use_depthwise': True,
-        'use_keras': True
-    },
-    {
-        'use_depthwise': False,
-        'use_keras': False
-    },
-    {
-        'use_depthwise': True,
-        'use_keras': False
+        'use_depthwise': True
     },
 )
 class SsdMobilenetV2FpnFeatureExtractorTest(
@@ -71,34 +63,20 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
     Returns:
       an ssd_meta_arch.SSDFeatureExtractor object.
     """
+    del use_keras
     min_depth = 32
-    if use_keras:
-      return (ssd_mobilenet_v2_fpn_keras_feature_extractor
-              .SSDMobileNetV2FpnKerasFeatureExtractor(
-                  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,
-                  use_explicit_padding=use_explicit_padding,
-                  use_depthwise=use_depthwise,
-                  name='MobilenetV2_FPN'))
-    else:
-      return (ssd_mobilenet_v2_fpn_feature_extractor
-              .SSDMobileNetV2FpnFeatureExtractor(
-                  is_training,
-                  depth_multiplier,
-                  min_depth,
-                  pad_to_multiple,
-                  self.conv_hyperparams_fn,
-                  use_depthwise=use_depthwise,
-                  use_explicit_padding=use_explicit_padding))
+    return (ssd_mobilenet_v2_fpn_feature_extractor
+            .SSDMobileNetV2FpnFeatureExtractor(
+                is_training,
+                depth_multiplier,
+                min_depth,
+                pad_to_multiple,
+                self.conv_hyperparams_fn,
+                use_depthwise=use_depthwise,
+                use_explicit_padding=use_explicit_padding))
 
-  def test_extract_features_returns_correct_shapes_256(self, use_keras,
-                                                       use_depthwise):
+  def test_extract_features_returns_correct_shapes_256(self, use_depthwise):
+    use_keras = False
     image_height = 256
     image_width = 256
     depth_multiplier = 1.0
@@ -127,8 +105,8 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
         use_keras=use_keras,
         use_depthwise=use_depthwise)
 
-  def test_extract_features_returns_correct_shapes_384(self, use_keras,
-                                                       use_depthwise):
+  def test_extract_features_returns_correct_shapes_384(self, use_depthwise):
+    use_keras = False
     image_height = 320
     image_width = 320
     depth_multiplier = 1.0
@@ -157,8 +135,9 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
         use_keras=use_keras,
         use_depthwise=use_depthwise)
 
-  def test_extract_features_with_dynamic_image_shape(self, use_keras,
+  def test_extract_features_with_dynamic_image_shape(self,
                                                      use_depthwise):
+    use_keras = False
     image_height = 256
     image_width = 256
     depth_multiplier = 1.0
@@ -188,7 +167,8 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
         use_depthwise=use_depthwise)
 
   def test_extract_features_returns_correct_shapes_with_pad_to_multiple(
-      self, use_keras, use_depthwise):
+      self, use_depthwise):
+    use_keras = False
     image_height = 299
     image_width = 299
     depth_multiplier = 1.0
@@ -218,7 +198,8 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
         use_depthwise=use_depthwise)
 
   def test_extract_features_returns_correct_shapes_enforcing_min_depth(
-      self, use_keras, use_depthwise):
+      self, use_depthwise):
+    use_keras = False
     image_height = 256
     image_width = 256
     depth_multiplier = 0.5**12
@@ -248,7 +229,8 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
         use_depthwise=use_depthwise)
 
   def test_extract_features_raises_error_with_invalid_image_size(
-      self, use_keras, use_depthwise):
+      self, use_depthwise):
+    use_keras = False
     image_height = 32
     image_width = 32
     depth_multiplier = 1.0
@@ -261,8 +243,9 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
         use_keras=use_keras,
         use_depthwise=use_depthwise)
 
-  def test_preprocess_returns_correct_value_range(self, use_keras,
+  def test_preprocess_returns_correct_value_range(self,
                                                   use_depthwise):
+    use_keras = False
     image_height = 256
     image_width = 256
     depth_multiplier = 1
@@ -276,7 +259,8 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
     preprocessed_image = feature_extractor.preprocess(test_image)
     self.assertTrue(np.all(np.less_equal(np.abs(preprocessed_image), 1.0)))
 
-  def test_variables_only_created_in_scope(self, use_keras, use_depthwise):
+  def test_variables_only_created_in_scope(self, use_depthwise):
+    use_keras = False
     depth_multiplier = 1
     pad_to_multiple = 1
     scope_name = 'MobilenetV2'
@@ -287,7 +271,8 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
         use_keras=use_keras,
         use_depthwise=use_depthwise)
 
-  def test_fused_batchnorm(self, use_keras, use_depthwise):
+  def test_fused_batchnorm(self, use_depthwise):
+    use_keras = False
     image_height = 256
     image_width = 256
     depth_multiplier = 1
@@ -300,15 +285,13 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
         use_keras=use_keras,
         use_depthwise=use_depthwise)
     preprocessed_image = feature_extractor.preprocess(image_placeholder)
-    if use_keras:
-      _ = feature_extractor(preprocessed_image)
-    else:
-      _ = feature_extractor.extract_features(preprocessed_image)
+    _ = feature_extractor.extract_features(preprocessed_image)
     self.assertTrue(
         any('FusedBatchNorm' in op.type
             for op in tf.get_default_graph().get_operations()))
 
-  def test_variable_count(self, use_keras, use_depthwise):
+  def test_variable_count(self, use_depthwise):
+    use_keras = False
     depth_multiplier = 1
     pad_to_multiple = 1
     variables = self.get_feature_extractor_variables(
@@ -321,8 +304,9 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
       expected_variables_len = 278
     self.assertEqual(len(variables), expected_variables_len)
 
-  def test_get_expected_feature_map_variable_names(self, use_keras,
+  def test_get_expected_feature_map_variable_names(self,
                                                    use_depthwise):
+    use_keras = False
     depth_multiplier = 1.0
     pad_to_multiple = 1
 
@@ -360,44 +344,6 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
         'MobilenetV2/fpn/projection_2/weights',
         'MobilenetV2/fpn/projection_3/weights',
     ])
-    keras_expected_feature_maps_variables = set([
-        # Keras Mobilenet V2 feature maps
-        'MobilenetV2_FPN/block_4_depthwise/depthwise_kernel',
-        'MobilenetV2_FPN/block_7_depthwise/depthwise_kernel',
-        'MobilenetV2_FPN/block_14_depthwise/depthwise_kernel',
-        'MobilenetV2_FPN/Conv_1/kernel',
-        # FPN layers
-        'MobilenetV2_FPN/bottom_up_Conv2d_20_conv/kernel',
-        'MobilenetV2_FPN/bottom_up_Conv2d_21_conv/kernel',
-        'MobilenetV2_FPN/FeatureMaps/top_down/smoothing_1_conv/kernel',
-        'MobilenetV2_FPN/FeatureMaps/top_down/smoothing_2_conv/kernel',
-        'MobilenetV2_FPN/FeatureMaps/top_down/projection_1/kernel',
-        'MobilenetV2_FPN/FeatureMaps/top_down/projection_2/kernel',
-        'MobilenetV2_FPN/FeatureMaps/top_down/projection_3/kernel'
-    ])
-    keras_expected_feature_maps_variables_with_depthwise = set([
-        # Keras Mobilenet V2 feature maps
-        'MobilenetV2_FPN/block_4_depthwise/depthwise_kernel',
-        'MobilenetV2_FPN/block_7_depthwise/depthwise_kernel',
-        'MobilenetV2_FPN/block_14_depthwise/depthwise_kernel',
-        'MobilenetV2_FPN/Conv_1/kernel',
-        # FPN layers
-        'MobilenetV2_FPN/bottom_up_Conv2d_20_depthwise_conv/depthwise_kernel',
-        'MobilenetV2_FPN/bottom_up_Conv2d_20_depthwise_conv/pointwise_kernel',
-        'MobilenetV2_FPN/bottom_up_Conv2d_21_depthwise_conv/depthwise_kernel',
-        'MobilenetV2_FPN/bottom_up_Conv2d_21_depthwise_conv/pointwise_kernel',
-        ('MobilenetV2_FPN/FeatureMaps/top_down/smoothing_1_depthwise_conv/'
-         'depthwise_kernel'),
-        ('MobilenetV2_FPN/FeatureMaps/top_down/smoothing_1_depthwise_conv/'
-         'pointwise_kernel'),
-        ('MobilenetV2_FPN/FeatureMaps/top_down/smoothing_2_depthwise_conv/'
-         'depthwise_kernel'),
-        ('MobilenetV2_FPN/FeatureMaps/top_down/smoothing_2_depthwise_conv/'
-         'pointwise_kernel'),
-        'MobilenetV2_FPN/FeatureMaps/top_down/projection_1/kernel',
-        'MobilenetV2_FPN/FeatureMaps/top_down/projection_2/kernel',
-        'MobilenetV2_FPN/FeatureMaps/top_down/projection_3/kernel'
-    ])
 
     g = tf.Graph()
     with g.as_default():
@@ -407,18 +353,12 @@ class SsdMobilenetV2FpnFeatureExtractorTest(
           pad_to_multiple,
           use_keras=use_keras,
           use_depthwise=use_depthwise)
-      if use_keras:
-        _ = feature_extractor(preprocessed_inputs)
-        expected_feature_maps_variables = keras_expected_feature_maps_variables
-        if use_depthwise:
-          expected_feature_maps_variables = (
-              keras_expected_feature_maps_variables_with_depthwise)
-      else:
-        _ = feature_extractor.extract_features(preprocessed_inputs)
-        expected_feature_maps_variables = slim_expected_feature_maps_variables
-        if use_depthwise:
-          expected_feature_maps_variables = (
-              slim_expected_feature_maps_variables_with_depthwise)
+
+      _ = feature_extractor.extract_features(preprocessed_inputs)
+      expected_feature_maps_variables = slim_expected_feature_maps_variables
+      if use_depthwise:
+        expected_feature_maps_variables = (
+            slim_expected_feature_maps_variables_with_depthwise)
       actual_variable_set = set([
           var.op.name for var in g.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
       ])

research/object_detection/models/ssd_mobilenet_v2_fpn_feature_extractor_tf2_test.py

+# Copyright 2018 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_mobilenet_v2_fpn_feature_extractor.
+
+By using parameterized test decorator, this test serves for both Slim-based and
+Keras-based Mobilenet V2 FPN feature extractors in SSD.
+"""
+import unittest
+from absl.testing import parameterized
+import numpy as np
+import tensorflow.compat.v1 as tf
+
+from object_detection.models import ssd_feature_extractor_test
+from object_detection.models import ssd_mobilenet_v2_fpn_keras_feature_extractor
+from object_detection.utils import tf_version
+
+
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
+@parameterized.parameters(
+    {
+        'use_depthwise': False,
+    },
+    {
+        'use_depthwise': True,
+    },
+)
+class SsdMobilenetV2FpnFeatureExtractorTest(
+    ssd_feature_extractor_test.SsdFeatureExtractorTestBase):
+
+  def _create_feature_extractor(self,
+                                depth_multiplier,
+                                pad_to_multiple,
+                                is_training=True,
+                                use_explicit_padding=False,
+                                use_keras=False,
+                                use_depthwise=False):
+    """Constructs a new feature extractor.
+
+    Args:
+      depth_multiplier: float depth multiplier for feature extractor
+      pad_to_multiple: the nearest multiple to zero pad the input height and
+        width dimensions to.
+      is_training: whether the network is in training mode.
+      use_explicit_padding: Use 'VALID' padding for convolutions, but prepad
+        inputs so that the output dimensions are the same as if 'SAME' padding
+        were used.
+      use_keras: if True builds a keras-based feature extractor, if False builds
+        a slim-based one.
+      use_depthwise: Whether to use depthwise convolutions.
+    Returns:
+      an ssd_meta_arch.SSDFeatureExtractor object.
+    """
+    del use_keras
+    min_depth = 32
+    return (ssd_mobilenet_v2_fpn_keras_feature_extractor
+            .SSDMobileNetV2FpnKerasFeatureExtractor(
+                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,
+                use_explicit_padding=use_explicit_padding,
+                use_depthwise=use_depthwise,
+                name='MobilenetV2_FPN'))
+
+  def test_extract_features_returns_correct_shapes_256(self,
+                                                       use_depthwise):
+    use_keras = True
+    image_height = 256
+    image_width = 256
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 32, 32, 256), (2, 16, 16, 256),
+                                  (2, 8, 8, 256), (2, 4, 4, 256),
+                                  (2, 2, 2, 256)]
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=False,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=True,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+
+  def test_extract_features_returns_correct_shapes_384(self,
+                                                       use_depthwise):
+    use_keras = True
+    image_height = 320
+    image_width = 320
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 40, 40, 256), (2, 20, 20, 256),
+                                  (2, 10, 10, 256), (2, 5, 5, 256),
+                                  (2, 3, 3, 256)]
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=False,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=True,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+
+  def test_extract_features_with_dynamic_image_shape(self,
+                                                     use_depthwise):
+    use_keras = True
+    image_height = 256
+    image_width = 256
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 32, 32, 256), (2, 16, 16, 256),
+                                  (2, 8, 8, 256), (2, 4, 4, 256),
+                                  (2, 2, 2, 256)]
+    self.check_extract_features_returns_correct_shapes_with_dynamic_inputs(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=False,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+    self.check_extract_features_returns_correct_shapes_with_dynamic_inputs(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=True,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+
+  def test_extract_features_returns_correct_shapes_with_pad_to_multiple(
+      self, use_depthwise):
+    use_keras = True
+    image_height = 299
+    image_width = 299
+    depth_multiplier = 1.0
+    pad_to_multiple = 32
+    expected_feature_map_shape = [(2, 40, 40, 256), (2, 20, 20, 256),
+                                  (2, 10, 10, 256), (2, 5, 5, 256),
+                                  (2, 3, 3, 256)]
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=False,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=True,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+
+  def test_extract_features_returns_correct_shapes_enforcing_min_depth(
+      self, use_depthwise):
+    use_keras = True
+    image_height = 256
+    image_width = 256
+    depth_multiplier = 0.5**12
+    pad_to_multiple = 1
+    expected_feature_map_shape = [(2, 32, 32, 32), (2, 16, 16, 32),
+                                  (2, 8, 8, 32), (2, 4, 4, 32),
+                                  (2, 2, 2, 32)]
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=False,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+    self.check_extract_features_returns_correct_shape(
+        2,
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        expected_feature_map_shape,
+        use_explicit_padding=True,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+
+  def test_extract_features_raises_error_with_invalid_image_size(
+      self, use_depthwise=False):
+    use_keras = True
+    image_height = 32
+    image_width = 32
+    depth_multiplier = 1.0
+    pad_to_multiple = 1
+    self.check_extract_features_raises_error_with_invalid_image_size(
+        image_height,
+        image_width,
+        depth_multiplier,
+        pad_to_multiple,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+
+  def test_preprocess_returns_correct_value_range(self,
+                                                  use_depthwise):
+    use_keras = True
+    image_height = 256
+    image_width = 256
+    depth_multiplier = 1
+    pad_to_multiple = 1
+    test_image = np.random.rand(2, image_height, image_width, 3)
+    feature_extractor = self._create_feature_extractor(
+        depth_multiplier,
+        pad_to_multiple,
+        use_keras=use_keras,
+        use_depthwise=use_depthwise)
+    preprocessed_image = feature_extractor.preprocess(test_image)
+    self.assertTrue(np.all(np.less_equal(np.abs(preprocessed_image), 1.0)))
+
+if __name__ == '__main__':
+  tf.test.main()

research/object_detection/models/ssd_mobilenet_v2_fpn_keras_feature_extractor.py

@@ -123,7 +123,7 @@ class SSDMobileNetV2FpnKerasFeatureExtractor(
       self._conv_defs = _create_modified_mobilenet_config()
     self._use_native_resize_op = use_native_resize_op
     self._feature_blocks = ['layer_4', 'layer_7', 'layer_14', 'layer_19']
-    self._mobilenet_v2 = None
+    self.classification_backbone = None
     self._fpn_features_generator = None
     self._coarse_feature_layers = []
 
@@ -147,7 +147,7 @@ class SSDMobileNetV2FpnKerasFeatureExtractor(
       outputs.append(full_mobilenet_v2.get_layer(output_layer_name).output)
     layer_19 = full_mobilenet_v2.get_layer(name='out_relu').output
     outputs.append(layer_19)
-    self._mobilenet_v2 = tf.keras.Model(
+    self.classification_backbone = tf.keras.Model(
         inputs=full_mobilenet_v2.inputs,
         outputs=outputs)
     # pylint:disable=g-long-lambda
@@ -216,7 +216,7 @@ class SSDMobileNetV2FpnKerasFeatureExtractor(
     preprocessed_inputs = shape_utils.check_min_image_dim(
         33, preprocessed_inputs)
 
-    image_features = self._mobilenet_v2(
+    image_features = self.classification_backbone(
         ops.pad_to_multiple(preprocessed_inputs, self._pad_to_multiple))
 
     feature_block_list = []
@@ -241,3 +241,14 @@ class SSDMobileNetV2FpnKerasFeatureExtractor(
         last_feature_map = layer(last_feature_map)
       feature_maps.append(last_feature_map)
     return feature_maps
+
+  def restore_from_classification_checkpoint_fn(self, feature_extractor_scope):
+    """Returns a map for restoring from an (object-based) checkpoint.
+
+    Args:
+      feature_extractor_scope: A scope name for the feature extractor (unused).
+
+    Returns:
+      A dict mapping keys to Keras models
+    """
+    return {'feature_extractor': self.classification_backbone}

research/object_detection/models/ssd_mobilenet_v2_keras_feature_extractor.py

@@ -97,7 +97,7 @@ class SSDMobileNetV2KerasFeatureExtractor(
         'use_explicit_padding': self._use_explicit_padding,
     }
 
-    self.mobilenet_v2 = None
+    self.classification_backbone = None
     self.feature_map_generator = None
 
   def build(self, input_shape):
@@ -114,7 +114,7 @@ class SSDMobileNetV2KerasFeatureExtractor(
     conv2d_11_pointwise = full_mobilenet_v2.get_layer(
         name='block_13_expand_relu').output
     conv2d_13_pointwise = full_mobilenet_v2.get_layer(name='out_relu').output
-    self.mobilenet_v2 = tf.keras.Model(
+    self.classification_backbone = tf.keras.Model(
         inputs=full_mobilenet_v2.inputs,
         outputs=[conv2d_11_pointwise, conv2d_13_pointwise])
     self.feature_map_generator = (
@@ -158,7 +158,7 @@ class SSDMobileNetV2KerasFeatureExtractor(
     preprocessed_inputs = shape_utils.check_min_image_dim(
         33, preprocessed_inputs)
 
-    image_features = self.mobilenet_v2(
+    image_features = self.classification_backbone(
         ops.pad_to_multiple(preprocessed_inputs, self._pad_to_multiple))
 
     feature_maps = self.feature_map_generator({
@@ -166,3 +166,14 @@ class SSDMobileNetV2KerasFeatureExtractor(
         'layer_19': image_features[1]})
 
     return list(feature_maps.values())
+
+  def restore_from_classification_checkpoint_fn(self, feature_extractor_scope):
+    """Returns a map for restoring from an (object-based) checkpoint.
+
+    Args:
+      feature_extractor_scope: A scope name for the feature extractor (unused).
+
+    Returns:
+      A dict mapping keys to Keras models
+    """
+    return {'feature_extractor': self.classification_backbone}

research/object_detection/models/ssd_mobilenet_v2_mnasfpn_feature_extractor_test.py → research/object_detection/models/ssd_mobilenet_v2_mnasfpn_feature_extractor_tf1_test.py

@@ -14,13 +14,16 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for ssd_mobilenet_v2_nas_fpn_feature_extractor."""
+import unittest
 import numpy as np
 import tensorflow.compat.v1 as tf
 
 from object_detection.models import ssd_feature_extractor_test
 from object_detection.models import ssd_mobilenet_v2_mnasfpn_feature_extractor as mnasfpn_feature_extractor
+from object_detection.utils import tf_version
 
 
+@unittest.skipIf(tf_version.is_tf2(), 'Skipping TF1.X only test.')
 class SsdMobilenetV2MnasFPNFeatureExtractorTest(
     ssd_feature_extractor_test.SsdFeatureExtractorTestBase):