Merge branch 'master' of https://github.com/tensorflow/models into run_superglue

official/vision/beta/modeling/factory_test.py

-# Lint as: python3
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +11,8 @@
 # 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.
-# ==============================================================================
+
+# Lint as: python3
 """Tests for factory.py."""
 
 # Import libraries

official/vision/beta/modeling/heads/__init__.py

+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# 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.
+
+# Lint as: python3
+"""Heads package definition."""
+
+from official.vision.beta.modeling.heads.dense_prediction_heads import RetinaNetHead
+from official.vision.beta.modeling.heads.dense_prediction_heads import RPNHead
+from official.vision.beta.modeling.heads.instance_heads import DetectionHead
+from official.vision.beta.modeling.heads.instance_heads import MaskHead
+from official.vision.beta.modeling.heads.segmentation_heads import SegmentationHead

official/vision/beta/modeling/heads/dense_prediction_heads.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,8 +11,8 @@
 # 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.
-# ==============================================================================
-"""Dense prediction heads."""
+
+"""Contains definitions of dense prediction heads."""
 
 # Import libraries
 import numpy as np
@@ -23,7 +23,7 @@ from official.modeling import tf_utils
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class RetinaNetHead(tf.keras.layers.Layer):
-  """RetinaNet head."""
+  """Creates a RetinaNet head."""
 
   def __init__(self,
                min_level,
@@ -40,31 +40,30 @@ class RetinaNetHead(tf.keras.layers.Layer):
                kernel_regularizer=None,
                bias_regularizer=None,
                **kwargs):
-    """Initialize params to build RetinaNet head.
+    """Initializes a RetinaNet head.
 
     Args:
-      min_level: `int` number of minimum feature level.
-      max_level: `int` number of maximum feature level.
-      num_classes: `int` number of classes to predict.
-      num_anchors_per_location: `int` number of number of anchors per pixel
+      min_level: An `int` number of minimum feature level.
+      max_level: An `int` number of maximum feature level.
+      num_classes: An `int` number of classes to predict.
+      num_anchors_per_location: An `int` number of number of anchors per pixel
         location.
-      num_convs: `int` number that represents the number of the intermediate
+      num_convs: An `int` number that represents the number of the intermediate
         conv layers before the prediction.
-      num_filters: `int` number that represents the number of filters of the
+      num_filters: An `int` number that represents the number of filters of the
         intermediate conv layers.
-      use_separable_conv: `bool`, indicating whether the separable conv layers
-        is used.
-      activation: `string`, indicating which activation is used, e.g. 'relu',
+      use_separable_conv: A `bool` that indicates whether the separable
+        convolution layers is used.
+      activation: A `str` that indicates which activation is used, e.g. 'relu',
         'swish', etc.
-      use_sync_bn: `bool`, whether to use synchronized batch normalization
-        across different replicas.
-      norm_momentum: `float`, the momentum parameter of the normalization
-        layers.
-      norm_epsilon: `float`, the epsilon parameter of the normalization layers.
-      kernel_regularizer: `tf.keras.regularizers.Regularizer` object for layer
-        kernal.
-      bias_regularizer: `tf.keras.regularizers.Regularizer` object for bias.
-      **kwargs: other keyword arguments passed to Layer.
+      use_sync_bn: A `bool` that indicates whether to use synchronized batch
+        normalization across different replicas.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default is None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
+      **kwargs: Additional keyword arguments to be passed.
     """
     super(RetinaNetHead, self).__init__(**kwargs)
     self._config_dict = {
@@ -209,21 +208,22 @@ class RetinaNetHead(tf.keras.layers.Layer):
     """Forward pass of the RetinaNet head.
 
     Args:
-      features: a dict of tensors
-        - key: `str`, the level of the multilevel features.
-        - values: `Tensor`, the feature map tensors, whose shape is
+      features: A `dict` of `tf.Tensor` where
+        - key: A `str` of the level of the multilevel features.
+        - values: A `tf.Tensor`, the feature map tensors, whose shape is
             [batch, height_l, width_l, channels].
 
     Returns:
-      scores: a dict of tensors which includes scores of the predictions.
-        - key: `str`, the level of the multilevel predictions.
-        - values: `Tensor`, the box scores predicted from a particular feature
-            level, whose shape is
+      scores: A `dict` of `tf.Tensor` which includes scores of the predictions.
+        - key: A `str` of the level of the multilevel predictions.
+        - values: A `tf.Tensor` of the box scores predicted from a particular
+            feature level, whose shape is
             [batch, height_l, width_l, num_classes * num_anchors_per_location].
-      boxes: a dict of tensors which includes coordinates of the predictions.
-        - key: `str`, the level of the multilevel predictions.
-        - values: `Tensor`, the box scores predicted from a particular feature
-            level, whose shape is
+      boxes: A `dict` of `tf.Tensor` which includes coordinates of the
+        predictions.
+        - key: A `str` of the level of the multilevel predictions.
+        - values: A `tf.Tensor` of the box scores predicted from a particular
+            feature level, whose shape is
             [batch, height_l, width_l, 4 * num_anchors_per_location].
     """
     scores = {}
@@ -260,7 +260,7 @@ class RetinaNetHead(tf.keras.layers.Layer):
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class RPNHead(tf.keras.layers.Layer):
-  """Region Proposal Network head."""
+  """Creates a Region Proposal Network (RPN) head."""
 
   def __init__(self,
                min_level,
@@ -276,29 +276,29 @@ class RPNHead(tf.keras.layers.Layer):
                kernel_regularizer=None,
                bias_regularizer=None,
                **kwargs):
-    """Initialize params to build Region Proposal Network head.
+    """Initializes a Region Proposal Network head.
 
     Args:
-      min_level: `int` number of minimum feature level.
-      max_level: `int` number of maximum feature level.
-      num_anchors_per_location: `int` number of number of anchors per pixel
+      min_level: An `int` number of minimum feature level.
+      max_level: An `int` number of maximum feature level.
+      num_anchors_per_location: An `int` number of number of anchors per pixel
         location.
-      num_convs: `int` number that represents the number of the intermediate
-        conv layers before the prediction.
-      num_filters: `int` number that represents the number of filters of the
-        intermediate conv layers.
-      use_separable_conv: `bool`, indicating whether the separable conv layers
-        is used.
-      activation: `string`, indicating which activation is used, e.g. 'relu',
+      num_convs: An `int` number that represents the number of the intermediate
+        convolution layers before the prediction.
+      num_filters: An `int` number that represents the number of filters of the
+        intermediate convolution layers.
+      use_separable_conv: A `bool` that indicates whether the separable
+        convolution layers is used.
+      activation: A `str` that indicates which activation is used, e.g. 'relu',
         'swish', etc.
-      use_sync_bn: `bool`, whether to use synchronized batch normalization
-        across different replicas.
-      norm_momentum: `float`, the momentum parameter of the normalizaton layers.
-      norm_epsilon: `float`, the epsilon parameter of the normalization layers.
-      kernel_regularizer: `tf.keras.regularizers.Regularizer` object for layer
-        kernel.
-      bias_regularizer: `tf.keras.regularizers.Regularizer` object for bias.
-      **kwargs: other keyword arguments passed to Layer.
+      use_sync_bn: A `bool` that indicates whether to use synchronized batch
+        normalization across different replicas.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default is None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
+      **kwargs: Additional keyword arguments to be passed.
     """
     super(RPNHead, self).__init__(**kwargs)
     self._config_dict = {
@@ -428,6 +428,27 @@ class RPNHead(tf.keras.layers.Layer):
     super(RPNHead, self).build(input_shape)
 
   def call(self, features):
+    """Forward pass of the RPN head.
+
+    Args:
+      features: A `dict` of `tf.Tensor` where
+        - key: A `str` of the level of the multilevel features.
+        - values: A `tf.Tensor`, the feature map tensors, whose shape is [batch,
+          height_l, width_l, channels].
+
+    Returns:
+      scores: A `dict` of `tf.Tensor` which includes scores of the predictions.
+        - key: A `str` of the level of the multilevel predictions.
+        - values: A `tf.Tensor` of the box scores predicted from a particular
+            feature level, whose shape is
+            [batch, height_l, width_l, num_classes * num_anchors_per_location].
+      boxes: A `dict` of `tf.Tensor` which includes coordinates of the
+        predictions.
+        - key: A `str` of the level of the multilevel predictions.
+        - values: A `tf.Tensor` of the box scores predicted from a particular
+            feature level, whose shape is
+            [batch, height_l, width_l, 4 * num_anchors_per_location].
+    """
     scores = {}
     boxes = {}
     for i, level in enumerate(

official/vision/beta/modeling/heads/dense_prediction_heads_test.py

-# Lint as: python3
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +11,8 @@
 # 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.
-# ==============================================================================
+
+# Lint as: python3
 """Tests for dense_prediction_heads.py."""
 
 # Import libraries

official/vision/beta/modeling/heads/instance_heads.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,8 +11,8 @@
 # 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.
-# ==============================================================================
-"""Instance prediction heads."""
+
+"""Contains definitions of instance prediction heads."""
 
 # Import libraries
 import tensorflow as tf
@@ -22,7 +22,7 @@ from official.modeling import tf_utils
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class DetectionHead(tf.keras.layers.Layer):
-  """Detection head."""
+  """Creates a detection head."""
 
   def __init__(self,
                num_classes,
@@ -38,31 +38,30 @@ class DetectionHead(tf.keras.layers.Layer):
                kernel_regularizer=None,
                bias_regularizer=None,
                **kwargs):
-    """Initialize params to build the detection head.
+    """Initializes a detection head.
 
     Args:
-      num_classes: a integer for the number of classes.
-      num_convs: `int` number that represents the number of the intermediate
-        conv layers before the FC layers.
-      num_filters: `int` number that represents the number of filters of the
-        intermediate conv layers.
-      use_separable_conv: `bool`, indicating whether the separable conv layers
-        is used.
-      num_fcs: `int` number that represents the number of FC layers before the
-        predictions.
-      fc_dims: `int` number that represents the number of dimension of the FC
+      num_classes: An `int` for the number of classes.
+      num_convs: An `int` number that represents the number of the intermediate
+        convolution layers before the FC layers.
+      num_filters: An `int` number that represents the number of filters of the
+        intermediate convolution layers.
+      use_separable_conv: A `bool` that indicates whether the separable
+        convolution layers is used.
+      num_fcs: An `int` number that represents the number of FC layers before
+        the predictions.
+      fc_dims: An `int` number that represents the number of dimension of the FC
         layers.
-      activation: `string`, indicating which activation is used, e.g. 'relu',
+      activation: A `str` that indicates which activation is used, e.g. 'relu',
         'swish', etc.
-      use_sync_bn: `bool`, whether to use synchronized batch normalization
-        across different replicas.
-      norm_momentum: `float`, the momentum parameter of the normalization
-        layers.
-      norm_epsilon: `float`, the epsilon parameter of the normalization layers.
-      kernel_regularizer: `tf.keras.regularizers.Regularizer` object for layer
-        kernel.
-      bias_regularizer: `tf.keras.regularizers.Regularizer` object for bias.
-      **kwargs: other keyword arguments passed to Layer.
+      use_sync_bn: A `bool` that indicates whether to use synchronized batch
+        normalization across different replicas.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default is None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
+      **kwargs: Additional keyword arguments to be passed.
     """
     super(DetectionHead, self).__init__(**kwargs)
     self._config_dict = {
@@ -165,18 +164,17 @@ class DetectionHead(tf.keras.layers.Layer):
     super(DetectionHead, self).build(input_shape)
 
   def call(self, inputs, training=None):
-    """Box and class branches for the Mask-RCNN model.
+    """Forward pass of box and class branches for the Mask-RCNN model.
 
     Args:
-      inputs: ROI features, a tensor of shape
-        [batch_size, num_instances, roi_height, roi_width, roi_channels],
-        representing the ROI features.
-      training: a boolean indicating whether it is in `training` mode.
+      inputs: A `tf.Tensor` of the shape [batch_size, num_instances, roi_height,
+        roi_width, roi_channels], representing the ROI features.
+      training: a `bool` indicating whether it is in `training` mode.
 
     Returns:
-      class_outputs: a tensor with a shape of
+      class_outputs: A `tf.Tensor` of the shape
         [batch_size, num_rois, num_classes], representing the class predictions.
-      box_outputs: a tensor with a shape of
+      box_outputs: A `tf.Tensor` of the shape
         [batch_size, num_rois, num_classes * 4], representing the box
         predictions.
     """
@@ -211,7 +209,7 @@ class DetectionHead(tf.keras.layers.Layer):
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class MaskHead(tf.keras.layers.Layer):
-  """Mask head."""
+  """Creates a mask head."""
 
   def __init__(self,
                num_classes,
@@ -227,31 +225,30 @@ class MaskHead(tf.keras.layers.Layer):
                bias_regularizer=None,
                class_agnostic=False,
                **kwargs):
-    """Initialize params to build the mask head.
+    """Initializes a mask head.
 
     Args:
-      num_classes: `int`, the number of classes.
-      upsample_factor: `int`, >= 1, the upsample factor to generate the
-        final predicted masks.
-      num_convs: `int` number that represents the number of the intermediate
-        conv layers before the mask prediction layers.
-      num_filters: `int` number that represents the number of filters of the
-        intermediate conv layers.
-      use_separable_conv: `bool`, indicating whether the separable conv layers
-        is used.
-      activation: `string`, indicating which activation is used, e.g. 'relu',
+      num_classes: An `int` of the number of classes.
+      upsample_factor: An `int` that indicates the upsample factor to generate
+        the final predicted masks. It should be >= 1.
+      num_convs: An `int` number that represents the number of the intermediate
+        convolution layers before the mask prediction layers.
+      num_filters: An `int` number that represents the number of filters of the
+        intermediate convolution layers.
+      use_separable_conv: A `bool` that indicates whether the separable
+        convolution layers is used.
+      activation: A `str` that indicates which activation is used, e.g. 'relu',
         'swish', etc.
-      use_sync_bn: `bool`, whether to use synchronized batch normalization
-        across different replicas.
-      norm_momentum: `float`, the momentum parameter of the normalization
-        layers.
-      norm_epsilon: `float`, the epsilon parameter of the normalization layers.
-      kernel_regularizer: `tf.keras.regularizers.Regularizer` object for layer
-        kernel.
-      bias_regularizer: `tf.keras.regularizers.Regularizer` object for bias.
-      class_agnostic: `bool`, if set, we use a single channel mask head that
+      use_sync_bn: A `bool` that indicates whether to use synchronized batch
+        normalization across different replicas.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default is None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
+      class_agnostic: A `bool`. If set, we use a single channel mask head that
         is shared between all classes.
-      **kwargs: other keyword arguments passed to Layer.
+      **kwargs: Additional keyword arguments to be passed.
     """
     super(MaskHead, self).__init__(**kwargs)
     self._config_dict = {
@@ -368,19 +365,18 @@ class MaskHead(tf.keras.layers.Layer):
     super(MaskHead, self).build(input_shape)
 
   def call(self, inputs, training=None):
-    """Mask branch for the Mask-RCNN model.
+    """Forward pass of mask branch for the Mask-RCNN model.
 
     Args:
-      inputs: a list of two tensors
-        inputs[0]: ROI features, a tensor of shape
-          [batch_size, num_instances, roi_height, roi_width, roi_channels],
-          representing the ROI features.
-        inputs[1]: ROI classes, a tensor of shape
-          [batch_size, num_instances], representing the classes of the ROIs.
-      training: a boolean indicating whether it is in `training` mode.
+      inputs: A `list` of two tensors where
+        inputs[0]: A `tf.Tensor` of shape [batch_size, num_instances,
+          roi_height, roi_width, roi_channels], representing the ROI features.
+        inputs[1]: A `tf.Tensor` of shape [batch_size, num_instances],
+          representing the classes of the ROIs.
+      training: A `bool` indicating whether it is in `training` mode.
 
     Returns:
-      mask_outputs: a tensor of shape
+      mask_outputs: A `tf.Tensor` of shape
         [batch_size, num_instances, roi_height * upsample_factor,
          roi_width * upsample_factor], representing the mask predictions.
     """

official/vision/beta/modeling/heads/instance_heads_test.py

-# Lint as: python3
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +11,8 @@
 # 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.
-# ==============================================================================
+
+# Lint as: python3
 """Tests for instance_heads.py."""
 
 # Import libraries

official/vision/beta/modeling/heads/segmentation_heads.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,8 +11,8 @@
 # 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.
-# ==============================================================================
-"""Segmentation heads."""
+
+"""Contains definitions of segmentation heads."""
 
 import tensorflow as tf
 
@@ -23,7 +23,7 @@ from official.vision.beta.ops import spatial_transform_ops
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class SegmentationHead(tf.keras.layers.Layer):
-  """Segmentation head."""
+  """Creates a segmentation head."""
 
   def __init__(self,
                num_classes,
@@ -41,38 +41,37 @@ class SegmentationHead(tf.keras.layers.Layer):
                kernel_regularizer=None,
                bias_regularizer=None,
                **kwargs):
-    """Initialize params to build segmentation head.
+    """Initializes a segmentation head.
 
     Args:
-      num_classes: `int` number of mask classification categories. The number of
-        classes does not include background class.
-      level: `int` or `str`, level to use to build segmentation head.
-      num_convs: `int` number of stacked convolution before the last prediction
-        layer.
-      num_filters: `int` number to specify the number of filters used.
+      num_classes: An `int` number of mask classification categories. The number
+        of classes does not include background class.
+      level: An `int` or `str`, level to use to build segmentation head.
+      num_convs: An `int` number of stacked convolution before the last
+        prediction layer.
+      num_filters: An `int` number to specify the number of filters used.
         Default is 256.
-      upsample_factor: `int` number to specify the upsampling factor to generate
-        finer mask. Default 1 means no upsampling is applied.
+      upsample_factor: An `int` number to specify the upsampling factor to
+        generate finer mask. Default 1 means no upsampling is applied.
       feature_fusion: One of `deeplabv3plus`, `pyramid_fusion`, or None. If
         `deeplabv3plus`, features from decoder_features[level] will be fused
         with low level feature maps from backbone. If `pyramid_fusion`,
         multiscale features will be resized and fused at the target level.
-      low_level: `int`, backbone level to be used for feature fusion. This arg
-        is used when feature_fusion is set to deeplabv3plus.
-      low_level_num_filters: `int`, reduced number of filters for the low
-        level features before fusing it with higher level features. This args is
-        only used when feature_fusion is set to deeplabv3plus.
-      activation: `string`, indicating which activation is used, e.g. 'relu',
+      low_level: An `int` of backbone level to be used for feature fusion. It is
+        used when feature_fusion is set to `deeplabv3plus`.
+      low_level_num_filters: An `int` of reduced number of filters for the low
+        level features before fusing it with higher level features. It is only
+        used when feature_fusion is set to `deeplabv3plus`.
+      activation: A `str` that indicates which activation is used, e.g. 'relu',
         'swish', etc.
-      use_sync_bn: `bool`, whether to use synchronized batch normalization
-        across different replicas.
-      norm_momentum: `float`, the momentum parameter of the normalization
-        layers.
-      norm_epsilon: `float`, the epsilon parameter of the normalization layers.
-      kernel_regularizer: `tf.keras.regularizers.Regularizer` object for layer
-        kernel.
-      bias_regularizer: `tf.keras.regularizers.Regularizer` object for bias.
-      **kwargs: other keyword arguments passed to Layer.
+      use_sync_bn: A `bool` that indicates whether to use synchronized batch
+        normalization across different replicas.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default is None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
+      **kwargs: Additional keyword arguments to be passed.
     """
     super(SegmentationHead, self).__init__(**kwargs)
 
@@ -160,17 +159,17 @@ class SegmentationHead(tf.keras.layers.Layer):
     """Forward pass of the segmentation head.
 
     Args:
-      backbone_output: a dict of tensors
-        - key: `str`, the level of the multilevel features.
-        - values: `Tensor`, the feature map tensors, whose shape is
+      backbone_output: A `dict` of tensors
+        - key: A `str` of the level of the multilevel features.
+        - values: A `tf.Tensor` of the feature map tensors, whose shape is
             [batch, height_l, width_l, channels].
-      decoder_output: a dict of tensors
-        - key: `str`, the level of the multilevel features.
-        - values: `Tensor`, the feature map tensors, whose shape is
+      decoder_output: A `dict` of tensors
+        - key: A `str` of the level of the multilevel features.
+        - values: A `tf.Tensor` of the feature map tensors, whose shape is
             [batch, height_l, width_l, channels].
     Returns:
-      segmentation prediction mask: `Tensor`, the segmentation mask scores
-        predicted from input feature.
+      segmentation prediction mask: A `tf.Tensor` of the segmentation mask
+        scores predicted from input features.
     """
     if self._config_dict['feature_fusion'] == 'deeplabv3plus':
       # deeplabv3+ feature fusion

official/vision/beta/modeling/heads/segmentation_heads_test.py

-# Lint as: python3
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +11,8 @@
 # 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.
-# ==============================================================================
+
+# Lint as: python3
 """Tests for segmentation_heads.py."""
 
 # Import libraries

official/vision/beta/modeling/layers/__init__.py

+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# 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.
+
+# Lint as: python3
+"""Layers package definition."""
+
+from official.vision.beta.modeling.layers.box_sampler import BoxSampler
+from official.vision.beta.modeling.layers.detection_generator import DetectionGenerator
+from official.vision.beta.modeling.layers.detection_generator import MultilevelDetectionGenerator
+from official.vision.beta.modeling.layers.mask_sampler import MaskSampler
+from official.vision.beta.modeling.layers.nn_blocks import BottleneckBlock
+from official.vision.beta.modeling.layers.nn_blocks import BottleneckResidualInner
+from official.vision.beta.modeling.layers.nn_blocks import DepthwiseSeparableConvBlock
+from official.vision.beta.modeling.layers.nn_blocks import InvertedBottleneckBlock
+from official.vision.beta.modeling.layers.nn_blocks import ResidualBlock
+from official.vision.beta.modeling.layers.nn_blocks import ResidualInner
+from official.vision.beta.modeling.layers.nn_blocks import ReversibleLayer
+from official.vision.beta.modeling.layers.nn_blocks_3d import BottleneckBlock3D
+from official.vision.beta.modeling.layers.nn_blocks_3d import SelfGating
+from official.vision.beta.modeling.layers.nn_layers import CausalConvMixin
+from official.vision.beta.modeling.layers.nn_layers import Conv2D
+from official.vision.beta.modeling.layers.nn_layers import Conv3D
+from official.vision.beta.modeling.layers.nn_layers import DepthwiseConv2D
+from official.vision.beta.modeling.layers.nn_layers import GlobalAveragePool3D
+from official.vision.beta.modeling.layers.nn_layers import PositionalEncoding
+from official.vision.beta.modeling.layers.nn_layers import Scale
+from official.vision.beta.modeling.layers.nn_layers import SpatialAveragePool3D
+from official.vision.beta.modeling.layers.nn_layers import SqueezeExcitation
+from official.vision.beta.modeling.layers.nn_layers import StochasticDepth
+from official.vision.beta.modeling.layers.nn_layers import TemporalSoftmaxPool
+from official.vision.beta.modeling.layers.roi_aligner import MultilevelROIAligner
+from official.vision.beta.modeling.layers.roi_generator import MultilevelROIGenerator
+from official.vision.beta.modeling.layers.roi_sampler import ROISampler

official/vision/beta/modeling/layers/box_sampler.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,7 +11,7 @@
 # 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.
-# ==============================================================================
+
 """Contains definitions of box sampler."""
 
 # Import libraries

official/vision/beta/modeling/layers/detection_generator.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,7 +11,7 @@
 # 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.
-# ==============================================================================
+
 """Contains definitions of generators to generate the final detections."""
 
 # Import libraries

official/vision/beta/modeling/layers/detection_generator_test.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,7 +11,7 @@
 # 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.
-# ==============================================================================
+
 """Tests for detection_generator.py."""
 # Import libraries
 

official/vision/beta/modeling/layers/mask_sampler.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,7 +11,7 @@
 # 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.
-# ==============================================================================
+
 """Contains definitions of mask sampler."""
 
 # Import libraries

official/vision/beta/modeling/layers/nn_blocks.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,7 +11,7 @@
 # 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.
-# ==============================================================================
+
 """Contains common building blocks for neural networks."""
 
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union, Text

official/vision/beta/modeling/layers/nn_blocks_3d.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,7 +11,7 @@
 # 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.
-# ==============================================================================
+
 """Contains common building blocks for 3D networks."""
 # Import libraries
 import tensorflow as tf
@@ -210,7 +210,7 @@ class BottleneckBlock3D(tf.keras.layers.Layer):
         'temporal_kernel_size': self._temporal_kernel_size,
         'temporal_strides': self._temporal_strides,
         'spatial_strides': self._spatial_strides,
-        'use_projection': self._use_projection,
+        'use_self_gating': self._use_self_gating,
         'kernel_initializer': self._kernel_initializer,
         'kernel_regularizer': self._kernel_regularizer,
         'bias_regularizer': self._bias_regularizer,

official/vision/beta/modeling/layers/nn_blocks_3d_test.py

-# Lint as: python3
-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +11,8 @@
 # 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.
-# ==============================================================================
+
+# Lint as: python3
 """Tests for resnet."""
 
 # Import libraries

official/vision/beta/modeling/layers/nn_blocks_test.py

-# Lint as: python3
-# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +11,8 @@
 # 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.
-# ==============================================================================
+
+# Lint as: python3
 """Tests for nn_blocks."""
 
 from typing import Any, Iterable, Tuple

official/vision/beta/modeling/layers/nn_layers.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,7 +11,7 @@
 # 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.
-# ==============================================================================
+
 """Contains common building blocks for neural networks."""
 
 from typing import Callable, Dict, List, Optional, Tuple, Union
@@ -232,7 +232,7 @@ class StochasticDepth(tf.keras.layers.Layer):
     batch_size = tf.shape(inputs)[0]
     random_tensor = keep_prob
     random_tensor += tf.random.uniform(
-        [batch_size, 1, 1, 1], dtype=inputs.dtype)
+        [batch_size] + [1] * (inputs.shape.rank - 1), dtype=inputs.dtype)
     binary_tensor = tf.floor(random_tensor)
     output = tf.math.divide(inputs, keep_prob) * binary_tensor
     return output
@@ -590,7 +590,7 @@ class GlobalAveragePool3D(tf.keras.layers.Layer):
       # regular global average pooling.
       # Shape: [batch_size, 1, 1, 1, channels]
       x = tf.reduce_sum(inputs, axis=(1, 2, 3), keepdims=True)
-      x = x / tf.cast(inputs.shape[2] * inputs.shape[3], x.dtype)
+      x = x / tf.cast(tf.shape(inputs)[2] * tf.shape(inputs)[3], x.dtype)
       x = x + buffer
 
       # Shape: [batch_size, 1, 1, 1, channels]
@@ -713,7 +713,7 @@ class CausalConvMixin:
     # When buffer padding, use 'valid' padding across time. The output shape
     # across time should be the input shape minus any padding, assuming
     # the stride across time is 1.
-    if self._use_buffered_input:
+    if self._use_buffered_input and spatial_output_shape[0] is not None:
       padding = self._compute_buffered_causal_padding(use_buffered_input=False)
       spatial_output_shape[0] -= sum(padding[1])
     return spatial_output_shape

official/vision/beta/modeling/layers/nn_layers_test.py

-# Lint as: python3
-# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,7 +11,8 @@
 # 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.
-# ==============================================================================
+
+# Lint as: python3
 """Tests for nn_layers."""
 
 # Import libraries

official/vision/beta/modeling/layers/roi_aligner.py

-# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,7 +11,7 @@
 # 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.
-# ==============================================================================
+
 """Contains definitions of ROI aligner."""
 
 import tensorflow as tf