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

Code before:
#
# 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.
# ==============================================================================
"""Library functions for ContextRCNN."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

# The negative value used in padding the invalid weights.
_NEGATIVE_PADDING_VALUE = -100000

KEY_NAME = 'key'
VALUE_NAME = 'val'
QUERY_NAME = 'query'
FEATURE_NAME = 'feature'

class ContextProjection(tf.keras.layers.Layer):
  """Custom layer to do batch normalization and projection."""
  def __init__(self, projection_dimension, freeze_batchnorm, **kwargs):
    self.batch_norm = tf.keras.layers.BatchNormalization(
        epsilon=0.001,
        center=True,
        scale=True,
        momentum=0.97,
        trainable=(not freeze_batchnorm))
    self.projection = tf.keras.layers.Dense(units=projection_dimension,
                                            activation=tf.nn.relu6,
                                            use_bias=True)
    super(ContextProjection, self).__init__(**kwargs)

  def build(self, input_shape):
    self.batch_norm.build(input_shape)
    self.projection.build(input_shape)

  def call(self, input_features, is_training=False):
    return self.projection(self.batch_norm(input_features, is_training))

class AttentionBlock(tf.keras.layers.Layer):
  """Custom layer to perform all attention."""
  def __init__(self, bottleneck_dimension, attention_temperature,
               freeze_batchnorm, output_dimension=None,
               is_training=False, **kwargs):
    self._key_proj = ContextProjection(bottleneck_dimension, freeze_batchnorm)
    self._val_proj = ContextProjection(bottleneck_dimension, freeze_batchnorm)
    self._query_proj = ContextProjection(bottleneck_dimension, freeze_batchnorm)

Review comment (performance):
`trainable` argument does not achieve the freezing behavior. Please use [FreezableBatchnorm](https://github.com/tensorflow/models/blob/master/research/object_detection/core/freezable_batch_norm.py) 

Code after:
#
# 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.
# ==============================================================================
"""Library functions for Context R-CNN."""
import tensorflow as tf

from object_detection.core import freezable_batch_norm

# The negative value used in padding the invalid weights.
_NEGATIVE_PADDING_VALUE = -100000

class ContextProjection(tf.keras.layers.Layer):
  """Custom layer to do batch normalization and projection."""
  def __init__(self, projection_dimension, **kwargs):
    self.batch_norm = freezable_batch_norm.FreezableBatchNorm(
        epsilon=0.001,
        center=True,
        scale=True,
        momentum=0.97,
        trainable=True)
    self.projection = tf.keras.layers.Dense(units=projection_dimension,
                                            activation=tf.nn.relu6,
                                            use_bias=True)
    super(ContextProjection, self).__init__(**kwargs)

  def build(self, input_shape):
    self.batch_norm.build(input_shape)
    self.projection.build(input_shape)

  def call(self, input_features, is_training=False):
    return self.projection(self.batch_norm(input_features, is_training))

class AttentionBlock(tf.keras.layers.Layer):
  """Custom layer to perform all attention."""
  def __init__(self, bottleneck_dimension, attention_temperature,
               output_dimension=None, is_training=False,
               name='AttentionBlock', **kwargs):
    """Constructs an attention block.

    Args:
      bottleneck_dimension: A int32 Tensor representing the bottleneck dimension
        for intermediate projections.
      attention_temperature: A float Tensor. It controls the temperature of the
        softmax for weights calculation. The formula for calculation as follows:
          weights = exp(weights / temperature) / sum(exp(weights / temperature))
      output_dimension: A int32 Tensor representing the last dimension of the
        output feature.