major fixes

research/object_detection/meta_architectures/context_rcnn_lib.py

@@ -218,8 +218,9 @@ def compute_box_context_attention(box_features, context_features,
                                     attention_temperature, valid_mask,
                                     is_training)
 
+  print("before expansion", output_features.shape)
   # Expands the dimension back to match with the original feature map.
   output_features = output_features[:, :, tf.newaxis, tf.newaxis, :]
-
+  print("after expansion", output_features.shape)
   return output_features
   
\ No newline at end of file

research/object_detection/meta_architectures/context_rcnn_lib_tf2_test.py

@@ -96,16 +96,9 @@ class ContextRcnnLibTest(parameterized.TestCase, test_case.TestCase,
   )
   def test_attention_block(self, bottleneck_dimension, output_dimension,
                            attention_temperature):
-    input_features = tf.ones([2, 3, 4], tf.float32)
-    context_features = tf.ones([2, 2, 3], tf.float32)
-    valid_mask = tf.constant([[True, True], [False, False]], tf.bool)
+    input_features = tf.ones([2, 8, 3, 3, 3], tf.float32)
+    context_features = tf.ones([2, 20, 10], tf.float32)
     is_training = False
-    #projection_layers = {context_rcnn_lib.KEY_NAME: context_rcnn_lib.ContextProjection(bottleneck_dimension, False), context_rcnn_lib.VALUE_NAME: context_rcnn_lib.ContextProjection(bottleneck_dimension, False),
-    #                     context_rcnn_lib.QUERY_NAME: context_rcnn_lib.ContextProjection(bottleneck_dimension, False)}
-    
-    #Add in the feature layer because this is further down the pipeline and it isn't automatically injected.
-    #projection_layers['feature'] = context_rcnn_lib.ContextProjection(output_dimension, False)
-
     attention_block = context_rcnn_lib.AttentionBlock(bottleneck_dimension, attention_temperature, False, output_dimension)
     valid_context_size = tf.random_uniform((2,),
                                              minval=0,
@@ -114,7 +107,7 @@ class ContextRcnnLibTest(parameterized.TestCase, test_case.TestCase,
     output_features = attention_block([input_features, context_features], is_training, valid_context_size)
 
     # Makes sure the shape is correct.
-    self.assertAllEqual(output_features.shape, [2, 3, output_dimension])
+    self.assertAllEqual(output_features.shape, [2, 8, 1, 1, output_dimension])
 
 if __name__ == '__main__':
   tf.test.main()

research/object_detection/meta_architectures/context_rcnn_lib_v2.py

@@ -74,6 +74,7 @@ class AttentionBlock(tf.keras.layers.Layer):
     print("CALLED")
     input_features, context_features = input_features
     print(input_features.shape)
+    print(context_features.shape)
 
     _, context_size, _ = context_features.shape
     valid_mask = compute_valid_mask(valid_context_size, context_size)

research/object_detection/meta_architectures/context_rcnn_meta_arch.py

@@ -349,7 +349,10 @@ class ContextRCNNMetaArch(faster_rcnn_meta_arch.FasterRCNNMetaArch):
     print(attention_features.shape)
 
     # Adds box features with attention features.
+    print("box", box_features.shape)
+    print("attention", attention_features.shape)
     box_features += attention_features
+    print("after adding", box_features.shape)
 
     flattened_feature_maps = self._flatten_first_two_dimensions(box_features)
 

research/object_detection/meta_architectures/context_rcnn_meta_arch_test.py

@@ -480,48 +480,6 @@ class ContextRCNNMetaArchTest(test_case.TestCase, parameterized.TestCase):
     _ = model.predict(preprocessed_inputs, true_image_shapes, **side_inputs)
     mock_context_rcnn_lib_v1.compute_box_context_attention.assert_called_once()
 
-  @unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
-  @mock.patch.object(context_rcnn_meta_arch, 'context_rcnn_lib_v2')
-  def test_prediction_mock_tf2(self, mock_context_rcnn_lib_v2):
-    """Mocks the context_rcnn_lib_v2 module to test the prediction.
-
-    Using mock object so that we can ensure compute_box_context_attention is
-    called in side the prediction function.
-
-    Args:
-      mock_context_rcnn_lib_v2: mock module for the context_rcnn_lib_v2.
-    """
-    model = self._build_model(
-        is_training=False,
-        number_of_stages=2,
-        second_stage_batch_size=6,
-        num_classes=42)
-    mock_tensor = tf.ones([2, 8, 3, 3, 3], tf.float32)
-
-    mock_context_rcnn_lib_v2.compute_box_context_attention.return_value = mock_tensor
-    print(mock_context_rcnn_lib_v2.compute_box_context_attention)
-    inputs_shape = (2, 20, 20, 3)
-    inputs = tf.cast(
-        tf.random_uniform(inputs_shape, minval=0, maxval=255, dtype=tf.int32),
-        dtype=tf.float32)
-    preprocessed_inputs, true_image_shapes = model.preprocess(inputs)
-    context_features = tf.random_uniform((2, 20, 10),
-                                         minval=0,
-                                         maxval=255,
-                                         dtype=tf.float32)
-    valid_context_size = tf.random_uniform((2,),
-                                           minval=0,
-                                           maxval=10,
-                                           dtype=tf.int32)
-    features = {
-        fields.InputDataFields.context_features: context_features,
-        fields.InputDataFields.valid_context_size: valid_context_size
-    }
-
-    side_inputs = model.get_side_inputs(features)
-    print("Predicting now")
-    _ = model.predict(preprocessed_inputs, true_image_shapes, **side_inputs)
-    #mock_context_rcnn_lib_v2.compute_box_context_attention.assert_called_once()
 
   @parameterized.named_parameters(
       {'testcase_name': 'static_shapes', 'static_shapes': True},
@@ -560,7 +518,8 @@ class ContextRCNNMetaArchTest(test_case.TestCase, parameterized.TestCase):
       }
 
       side_inputs = model.get_side_inputs(features)
-
+      print('preprocessed', preprocessed_inputs.shape)
+      print('context', context_features.shape)
       prediction_dict = model.predict(preprocessed_inputs, true_image_shapes,
                                       **side_inputs)
       return (prediction_dict['rpn_box_predictor_features'],