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

Code before:
    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)
    output_features = project_features(
        features, self._output_dimension, self._is_training,
        self._feature_proj, normalize=False)

    output_features = output_features[:, :, tf.newaxis, tf.newaxis, :]

    return output_features


Review comment (refactor):
looks like name_scope doesn't affect variable names. please remove this and add a constructor argument `name='AttentionBlock' to class AttentionBlock and pass it to the base class.

Code after:
    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])

    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,