Repository: tensorflow/models
File: research/object_detection/meta_architectures/context_rcnn_lib_tf2.py (Python)

Code before:
    self._val_proj = ContextProjection(bottleneck_dimension, freeze_batchnorm)
    self._query_proj = ContextProjection(bottleneck_dimension, freeze_batchnorm)
    self._feature_proj = None
    self._attention_temperature = attention_temperature
    self._freeze_batchnorm = freeze_batchnorm
    self._bottleneck_dimension = bottleneck_dimension
    self._output_dimension = output_dimension
    self._is_training = is_training
    super(AttentionBlock, self).__init__(**kwargs)

  def set_output_dimension(self, output_dim):
    self._output_dimension = output_dim

  def build(self, input_shapes):
    pass

  def call(self, box_features, context_features, valid_context_size):
    """Handles a call by performing attention."""
    _, context_size, _ = context_features.shape
    valid_mask = compute_valid_mask(valid_context_size, context_size)
    channels = box_features.shape[-1]

    #Build the feature projection layer
    if not self._output_dimension:
      self._output_dimension = channels
    if not self._feature_proj:
      self._feature_proj = ContextProjection(self._output_dimension,
                                             self._freeze_batchnorm)

    # Average pools over height and width dimension so that the shape of
    # box_features becomes [batch_size, max_num_proposals, channels].
    box_features = tf.reduce_mean(box_features, [2, 3])

    with tf.name_scope("AttentionBlock"):
      queries = project_features(
          box_features, self._bottleneck_dimension, self._is_training,
          self._query_proj, normalize=True)
      keys = project_features(
          context_features, self._bottleneck_dimension, self._is_training,
          self._key_proj, normalize=True)
      values = project_features(
          context_features, self._bottleneck_dimension, self._is_training,
          self._val_proj, normalize=True)

    weights = tf.matmul(queries, keys, transpose_b=True)

    weights, values = filter_weight_value(weights, values, valid_mask)

    weights = tf.nn.softmax(weights / self._attention_temperature)

    features = tf.matmul(weights, values)

Review comment (suggestion):
 self._feature_proj is always None. I think you mean to create it in the constructor if output_dimension is specified ? 

Also, when output dimension is unspecified, you should create this layer in `build` function that gets a list of `TensorShape`s corresponding to `box_features`, `context_features` and `valid_context_size`. See https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#build

Code after:
      output_dimension: A int32 Tensor representing the last dimension of the
        output feature.
      is_training: A boolean Tensor (affecting batch normalization).
      name: A string describing what to name the variables in this block.
      **kwargs: Additional keyword arguments.
    """

    self._key_proj = ContextProjection(bottleneck_dimension)
    self._val_proj = ContextProjection(bottleneck_dimension)
    self._query_proj = ContextProjection(bottleneck_dimension)
    self._feature_proj = None
    self._attention_temperature = attention_temperature
    self._bottleneck_dimension = bottleneck_dimension
    self._is_training = is_training
    self._output_dimension = output_dimension
    if self._output_dimension:
      self._feature_proj = ContextProjection(self._output_dimension)
    super(AttentionBlock, self).__init__(name=name, **kwargs)

  def build(self, input_shapes):
    """Finishes building the attention block.
    
    Args:
      input_shapes: the shape of the primary input box features.
    """
    if not self._feature_proj:
      self._output_dimension = input_shapes[-1]
      self._feature_proj = ContextProjection(self._output_dimension)

  def call(self, box_features, context_features, valid_context_size):
    """Handles a call by performing attention.

    Args:
      box_features: A float Tensor of shape [batch_size, input_size,
        num_input_features].
      context_features: A float Tensor of shape [batch_size, context_size,
        num_context_features].
      valid_context_size: A int32 Tensor of shape [batch_size].

    Returns:
      A float Tensor with shape [batch_size, input_size, num_input_features]
      containing output features after attention with context features.
    """

    _, context_size, _ = context_features.shape
    valid_mask = compute_valid_mask(valid_context_size, context_size)

    # Average pools over height and width dimension so that the shape of
    # box_features becomes [batch_size, max_num_proposals, channels].
    box_features = tf.reduce_mean(box_features, [2, 3])