reshape output for second stage

research/object_detection/models/faster_rcnn_resnet_v1_fpn_keras_feature_extractor.py

@@ -265,7 +265,8 @@ class FasterRCNNResnetV1FpnKerasFeatureExtractor(
         feature_extractor_model = tf.keras.models.Sequential([
             tf.keras.layers.Flatten(),
             tf.keras.layers.Dense(units=1024, activation='relu'),
-            tf.keras.layers.Dense(units=1024, activation='relu')
+            tf.keras.layers.Dense(units=1024, activation='relu'),
+            tf.keras.layers.Reshape((1, 1, 1024))
         ])
         return feature_extractor_model
 

research/object_detection/models/faster_rcnn_resnet_v1_fpn_keras_feature_extractor_tf2_test.py

@@ -91,4 +91,4 @@ class FasterRCNNResnetV1FpnKerasFeatureExtractorTest(tf.test.TestCase):
         model(proposal_feature_maps))
     features_shape = tf.shape(proposal_classifier_features)
 
-    self.assertAllEqual(features_shape.numpy(), [3, 1024])
+    self.assertAllEqual(features_shape.numpy(), [3, 1, 1, 1024])