Internal change

PiperOrigin-RevId: 361199986

official/vision/beta/modeling/backbones/efficientnet.py

@@ -51,14 +51,14 @@ SCALING_MAP = {
 
 
 def round_repeats(repeats, multiplier, skip=False):
-  """Round number of filters based on depth multiplier."""
+  """Returns rounded number of filters based on depth multiplier."""
   if skip or not multiplier:
     return repeats
   return int(math.ceil(multiplier * repeats))
 
 
 def block_spec_decoder(specs, width_scale, depth_scale):
-  """Decode specs for a block."""
+  """Decodes and returns specs for a block."""
   decoded_specs = []
   for s in specs:
     s = s + (
@@ -87,7 +87,13 @@ class BlockSpec(object):
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class EfficientNet(tf.keras.Model):
-  """Class to build EfficientNet family model."""
+  """Creates an EfficientNet family model.
+
+  This implements the EfficientNet model from:
+    Mingxing Tan, Quoc V. Le.
+    EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks.
+    (https://arxiv.org/pdf/1905.11946)
+  """
 
   def __init__(self,
                model_id,
@@ -102,25 +108,25 @@ class EfficientNet(tf.keras.Model):
                norm_momentum=0.99,
                norm_epsilon=0.001,
                **kwargs):
-    """EfficientNet initialization function.
+    """Initializes an EfficientNet model.
 
     Args:
-      model_id: `str` model id of EfficientNet.
-      input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
-      se_ratio: `float` squeeze and excitation ratio for inverted bottleneck
-        blocks.
-      stochastic_depth_drop_rate: `float` drop rate for drop connect layer.
-      kernel_initializer: kernel_initializer for convolutional layers.
-      kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
-        Default to None.
-      bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
+      model_id: A `str` of model ID of EfficientNet.
+      input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
+      se_ratio: A `float` of squeeze and excitation ratio for inverted
+        bottleneck blocks.
+      stochastic_depth_drop_rate: A `float` of drop rate for drop connect layer.
+      kernel_initializer: A `str` for kernel initializer of convolutional
+        layers.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default to None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
         Default to None.
-      activation: `str` name of the activation function.
-      use_sync_bn: if True, use synchronized batch normalization.
-      norm_momentum: `float` normalization omentum for the moving average.
-      norm_epsilon: `float` small float added to variance to avoid dividing by
-        zero.
-      **kwargs: keyword arguments to be passed.
+      activation: A `str` of name of the activation function.
+      use_sync_bn: If True, use synchronized batch normalization.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      **kwargs: Additional keyword arguments to be passed.
     """
     self._model_id = model_id
     self._input_specs = input_specs
@@ -203,12 +209,12 @@ class EfficientNet(tf.keras.Model):
     """Creates one group of blocks for the EfficientNet model.
 
     Args:
-      inputs: `Tensor` of size `[batch, channels, height, width]`.
-      specs: specifications for one inverted bottleneck block group.
-      name: `str`name for the block.
+      inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
+      specs: The specifications for one inverted bottleneck block group.
+      name: A `str` name for the block.
 
     Returns:
-      The output `Tensor` of the block layer.
+      The output `tf.Tensor` of the block layer.
     """
     if specs.block_fn == 'mbconv':
       block_fn = nn_blocks.InvertedBottleneckBlock
@@ -282,7 +288,7 @@ def build_efficientnet(
     input_specs: tf.keras.layers.InputSpec,
     model_config,
     l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
-  """Builds ResNet 3d backbone from a config."""
+  """Builds EfficientNet backbone from a config."""
   backbone_type = model_config.backbone.type
   backbone_cfg = model_config.backbone.get()
   norm_activation_config = model_config.norm_activation

official/vision/beta/modeling/backbones/factory.py

@@ -70,10 +70,10 @@ def register_backbone_builder(key: str):
   ```
 
   Args:
-    key: the key to look up the builder.
+    key: A `str` of key to look up the builder.
 
   Returns:
-    A callable for use as class decorator that registers the decorated class
+    A callable for using as class decorator that registers the decorated class
     for creation from an instance of task_config_cls.
   """
   return registry.register(_REGISTERED_BACKBONE_CLS, key)
@@ -85,12 +85,13 @@ def build_backbone(input_specs: tf.keras.layers.InputSpec,
   """Builds backbone from a config.
 
   Args:
-    input_specs: tf.keras.layers.InputSpec.
-    model_config: a OneOfConfig. Model config.
-    l2_regularizer: tf.keras.regularizers.Regularizer instance. Default to None.
+    input_specs: A `tf.keras.layers.InputSpec` of input.
+    model_config: A `OneOfConfig` of model config.
+    l2_regularizer: A `tf.keras.regularizers.Regularizer` object. Default to
+      None.
 
   Returns:
-    tf.keras.Model instance of the backbone.
+    A `tf.keras.Model` instance of the backbone.
   """
   backbone_builder = registry.lookup(_REGISTERED_BACKBONE_CLS,
                                      model_config.backbone.type)

official/vision/beta/modeling/backbones/mobilenet.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Contains definitions of Mobilenet Networks."""
+"""Contains definitions of MobileNet Networks."""
 
-from typing import Text, Optional, Dict, Any, Tuple
+from typing import Optional, Dict, Any, Tuple
 
 # Import libraries
 import dataclasses
@@ -41,8 +41,8 @@ class Conv2DBNBlock(tf.keras.layers.Layer):
       kernel_size: int = 3,
       strides: int = 1,
       use_bias: bool = False,
-      activation: Text = 'relu6',
-      kernel_initializer: Text = 'VarianceScaling',
+      activation: str = 'relu6',
+      kernel_initializer: str = 'VarianceScaling',
       kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
       bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
       use_normalization: bool = True,
@@ -53,25 +53,25 @@ class Conv2DBNBlock(tf.keras.layers.Layer):
     """A convolution block with batch normalization.
 
     Args:
-      filters: `int` number of filters for the first two convolutions. Note that
-        the third and final convolution will use 4 times as many filters.
-      kernel_size: `int` an integer specifying the height and width of the
-        2D convolution window.
-      strides: `int` block stride. If greater than 1, this block will ultimately
-        downsample the input.
-      use_bias: if True, use biase in the convolution layer.
-      activation: `str` name of the activation function.
-      kernel_initializer: kernel_initializer for convolutional layers.
-      kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
-                          Default to None.
-      bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
-                        Default to None.
-      use_normalization: if True, use batch normalization.
-      use_sync_bn: if True, use synchronized batch normalization.
-      norm_momentum: `float` normalization momentum for the moving average.
-      norm_epsilon: `float` small float added to variance to avoid dividing by
-        zero.
-      **kwargs: keyword arguments to be passed.
+      filters: An `int` number of filters for the first two convolutions. Note
+        that the third and final convolution will use 4 times as many filters.
+      kernel_size: An `int` specifying the height and width of the 2D
+        convolution window.
+      strides: An `int` of block stride. If greater than 1, this block will
+        ultimately downsample the input.
+      use_bias: If True, use bias in the convolution layer.
+      activation: A `str` name of the activation function.
+      kernel_initializer: A `str` for kernel initializer of convolutional
+        layers.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default to None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
+        Default to None.
+      use_normalization: If True, use batch normalization.
+      use_sync_bn: If True, use synchronized batch normalization.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      **kwargs: Additional keyword arguments to be passed.
     """
     super(Conv2DBNBlock, self).__init__(**kwargs)
     self._filters = filters
@@ -375,13 +375,13 @@ SUPPORTED_SPECS_MAP = {
 class BlockSpec(hyperparams.Config):
   """A container class that specifies the block configuration for MobileNet."""
 
-  block_fn: Text = 'convbn'
+  block_fn: str = 'convbn'
   kernel_size: int = 3
   strides: int = 1
   filters: int = 32
   use_bias: bool = False
   use_normalization: bool = True
-  activation: Text = 'relu6'
+  activation: str = 'relu6'
   # used for block type InvertedResConv
   expand_ratio: Optional[float] = 6.
   # used for block type InvertedResConv with SE
@@ -395,22 +395,22 @@ def block_spec_decoder(specs: Dict[Any, Any],
                        # set to 1 for mobilenetv1
                        divisible_by: int = 8,
                        finegrain_classification_mode: bool = True):
-  """Decode specs for a block.
+  """Decodes specs for a block.
 
   Args:
-    specs: `dict` specification of block specs of a mobilenet version.
-    filter_size_scale: `float` multiplier for the filter size
-      for all convolution ops. The value must be greater than zero. Typical
-      usage will be to set this value in (0, 1) to reduce the number of
-      parameters or computation cost of the model.
-    divisible_by: `int` ensures all inner dimensions are divisible by
+    specs: A `dict` specification of block specs of a mobilenet version.
+    filter_size_scale: A `float` multiplier for the filter size for all
+      convolution ops. The value must be greater than zero. Typical usage will
+      be to set this value in (0, 1) to reduce the number of parameters or
+      computation cost of the model.
+    divisible_by: An `int` that ensures all inner dimensions are divisible by
       this number.
-    finegrain_classification_mode: if True, the model
-      will keep the last layer large even for small multipliers. Following
-      https://arxiv.org/abs/1801.04381
+    finegrain_classification_mode: If True, the model will keep the last layer
+      large even for small multipliers, following
+      https://arxiv.org/abs/1801.04381.
 
   Returns:
-    List[BlockSpec]` defines structure of the base network.
+    A list of `BlockSpec` that defines structure of the base network.
   """
 
   spec_name = specs['spec_name']
@@ -449,66 +449,68 @@ def block_spec_decoder(specs: Dict[Any, Any],
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class MobileNet(tf.keras.Model):
-  """Class to build MobileNet family model."""
-
-  def __init__(self,
-               model_id: Text = 'MobileNetV2',
-               filter_size_scale: float = 1.0,
-               input_specs: layers.InputSpec = layers.InputSpec(
-                   shape=[None, None, None, 3]),
-               # The followings are for hyper-parameter tuning
-               norm_momentum: float = 0.99,
-               norm_epsilon: float = 0.001,
-               kernel_initializer: Text = 'VarianceScaling',
-               kernel_regularizer: Optional[regularizers.Regularizer] = None,
-               bias_regularizer: Optional[regularizers.Regularizer] = None,
-               # The followings should be kept the same most of the times
-               output_stride: int = None,
-               min_depth: int = 8,
-               # divisible is not used in MobileNetV1
-               divisible_by: int = 8,
-               stochastic_depth_drop_rate: float = 0.0,
-               regularize_depthwise: bool = False,
-               use_sync_bn: bool = False,
-               # finegrain is not used in MobileNetV1
-               finegrain_classification_mode: bool = True,
-               **kwargs):
-    """MobileNet initializer.
+  """Creates a MobileNet family model."""
+
+  def __init__(
+      self,
+      model_id: str = 'MobileNetV2',
+      filter_size_scale: float = 1.0,
+      input_specs: layers.InputSpec = layers.InputSpec(
+          shape=[None, None, None, 3]),
+      # The followings are for hyper-parameter tuning
+      norm_momentum: float = 0.99,
+      norm_epsilon: float = 0.001,
+      kernel_initializer: str = 'VarianceScaling',
+      kernel_regularizer: Optional[regularizers.Regularizer] = None,
+      bias_regularizer: Optional[regularizers.Regularizer] = None,
+      # The followings should be kept the same most of the times
+      output_stride: int = None,
+      min_depth: int = 8,
+      # divisible is not used in MobileNetV1
+      divisible_by: int = 8,
+      stochastic_depth_drop_rate: float = 0.0,
+      regularize_depthwise: bool = False,
+      use_sync_bn: bool = False,
+      # finegrain is not used in MobileNetV1
+      finegrain_classification_mode: bool = True,
+      **kwargs):
+    """Initializes a MobileNet model.
 
     Args:
-      model_id: `str` version of MobileNet. The supported values are
-       'MobileNetV1', 'MobileNetV2', 'MobileNetV3Large', 'MobileNetV3Small',
-        and 'MobileNetV3EdgeTPU'.
-      filter_size_scale: `float` multiplier for the filters (number of channels)
-        for all convolution ops. The value must be greater than zero. Typical
-        usage will be to set this value in (0, 1) to reduce the number of
-        parameters or computation cost of the model.
-      input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
-      norm_momentum: `float` normalization omentum for the moving average.
-      norm_epsilon: `float` small float added to variance to avoid dividing by
-        zero.
-      kernel_initializer: `str` kernel_initializer for convolutional layers.
-      kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
-        Default to None.
-      bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
+      model_id: A `str` of MobileNet version. The supported values are
+        `MobileNetV1`, `MobileNetV2`, `MobileNetV3Large`, `MobileNetV3Small`,
+        and `MobileNetV3EdgeTPU`.
+      filter_size_scale: A `float` of multiplier for the filters (number of
+        channels) for all convolution ops. The value must be greater than zero.
+        Typical usage will be to set this value in (0, 1) to reduce the number
+        of parameters or computation cost of the model.
+      input_specs: A `tf.keras.layers.InputSpec` of specs of the input tensor.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      kernel_initializer: A `str` for kernel initializer of convolutional
+        layers.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default to None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
         Default to None.
-      output_stride: `int` specifies the requested ratio of input to output
-        spatial resolution. If not None, then we invoke atrous convolution
-        if necessary to prevent the network from reducing the spatial resolution
-        of activation maps. Allowed values are 8 (accurate fully convolutional
-        mode), 16 (fast fully convolutional mode), 32 (classification mode).
-      min_depth: `int` minimum depth (number of channels) for all conv ops.
-        Enforced when filter_size_scale < 1, and not an active constraint when
-        filter_size_scale >= 1.
-      divisible_by: `int` ensures all inner dimensions are divisible by
+      output_stride: An `int` that specifies the requested ratio of input to
+        output spatial resolution. If not None, then we invoke atrous
+        convolution if necessary to prevent the network from reducing the
+        spatial resolution of activation maps. Allowed values are 8 (accurate
+        fully convolutional mode), 16 (fast fully convolutional mode), 32
+        (classification mode).
+      min_depth: An `int` of minimum depth (number of channels) for all
+        convolution ops. Enforced when filter_size_scale < 1, and not an active
+        constraint when filter_size_scale >= 1.
+      divisible_by: An `int` that ensures all inner dimensions are divisible by
         this number.
-      stochastic_depth_drop_rate: `float` drop rate for drop connect layer.
-      regularize_depthwise: if Ture, apply regularization on depthwise.
-      use_sync_bn: if True, use synchronized batch normalization.
-      finegrain_classification_mode: if True, the model
-        will keep the last layer large even for small multipliers. Following
-        https://arxiv.org/abs/1801.04381
-      **kwargs: keyword arguments to be passed.
+      stochastic_depth_drop_rate: A `float` of drop rate for drop connect layer.
+      regularize_depthwise: If Ture, apply regularization on depthwise.
+      use_sync_bn: If True, use synchronized batch normalization.
+      finegrain_classification_mode: If True, the model will keep the last layer
+        large even for small multipliers, following
+        https://arxiv.org/abs/1801.04381.
+      **kwargs: Additional keyword arguments to be passed.
     """
     if model_id not in SUPPORTED_SPECS_MAP:
       raise ValueError('The MobileNet version {} '
@@ -567,10 +569,10 @@ class MobileNet(tf.keras.Model):
   def _mobilenet_base(self,
                       inputs: tf.Tensor
                       ) -> Tuple[tf.Tensor, Dict[int, tf.Tensor]]:
-    """Build the base MobileNet architecture.
+    """Builds the base MobileNet architecture.
 
     Args:
-      inputs: Input tensor of shape [batch_size, height, width, channels].
+      inputs: A `tf.Tensor` of shape `[batch_size, height, width, channels]`.
 
     Returns:
       A tuple of output Tensor and dictionary that collects endpoints.
@@ -725,7 +727,7 @@ def build_mobilenet(
     input_specs: tf.keras.layers.InputSpec,
     model_config,
     l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
-  """Builds MobileNet 3d backbone from a config."""
+  """Builds MobileNet backbone from a config."""
   backbone_type = model_config.backbone.type
   backbone_cfg = model_config.backbone.get()
   norm_activation_config = model_config.norm_activation

official/vision/beta/modeling/backbones/resnet.py

@@ -12,12 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Contains definitions of Residual Networks.
-
-Residual networks (ResNets) were proposed in:
-[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
-    Deep Residual Learning for Image Recognition. arXiv:1512.03385
-"""
+"""Contains definitions of Residual Networks."""
 
 # Import libraries
 import tensorflow as tf
@@ -92,7 +87,13 @@ RESNET_SPECS = {
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class ResNet(tf.keras.Model):
-  """Class to build ResNet family model."""
+  """Creates a ResNet family model.
+
+  This implements the Deep Residual Network from:
+    Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun.
+    Deep Residual Learning for Image Recognition.
+    (https://arxiv.org/pdf/1512.03385)
+  """
 
   def __init__(self,
                model_id,
@@ -111,32 +112,31 @@ class ResNet(tf.keras.Model):
                kernel_regularizer=None,
                bias_regularizer=None,
                **kwargs):
-    """ResNet initialization function.
+    """Initializes a ResNet model.
 
     Args:
-      model_id: `int` depth of ResNet backbone model.
-      input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
-      depth_multiplier: `float` a depth multiplier to uniformaly scale up all
-        layers in channel size in ResNet.
-      stem_type: `str` stem type of ResNet. Default to `v0`. If set to `v1`,
-        use ResNet-D type stem (https://arxiv.org/abs/1812.01187).
-      resnetd_shortcut: `bool` whether to use ResNet-D shortcut in downsampling
-        blocks.
-      replace_stem_max_pool: `bool` if True, replace the max pool in stem with
-        a stride-2 conv,
-      se_ratio: `float` or None. Ratio of the Squeeze-and-Excitation layer.
-      init_stochastic_depth_rate: `float` initial stochastic depth rate.
-      activation: `str` name of the activation function.
-      use_sync_bn: if True, use synchronized batch normalization.
-      norm_momentum: `float` normalization omentum for the moving average.
-      norm_epsilon: `float` small float added to variance to avoid dividing by
-        zero.
-      kernel_initializer: kernel_initializer for convolutional layers.
-      kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
-                          Default to None.
-      bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
-                        Default to None.
-      **kwargs: keyword arguments to be passed.
+      model_id: An `int` of the depth of ResNet backbone model.
+      input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
+      depth_multiplier: A `float` of the depth multiplier to uniformaly scale up
+        all layers in channel size in ResNet.
+      stem_type: A `str` of stem type of ResNet. Default to `v0`. If set to
+        `v1`, use ResNet-D type stem (https://arxiv.org/abs/1812.01187).
+      resnetd_shortcut: A `bool` of whether to use ResNet-D shortcut in
+        downsampling blocks.
+      replace_stem_max_pool: A `bool` of whether to replace the max pool in stem
+        with a stride-2 conv,
+      se_ratio: A `float` or None. Ratio of the Squeeze-and-Excitation layer.
+      init_stochastic_depth_rate: A `float` of initial stochastic depth rate.
+      activation: A `str` name of the activation function.
+      use_sync_bn: If True, use synchronized batch normalization.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A small `float` added to variance to avoid dividing by zero.
+      kernel_initializer: A str for kernel initializer of convolutional layers.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default to None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
+        Default to None.
+      **kwargs: Additional keyword arguments to be passed.
     """
     self._model_id = model_id
     self._input_specs = input_specs
@@ -279,17 +279,20 @@ class ResNet(tf.keras.Model):
     """Creates one group of blocks for the ResNet model.
 
     Args:
-      inputs: `Tensor` of size `[batch, channels, height, width]`.
-      filters: `int` number of filters for the first convolution of the layer.
-      strides: `int` stride to use for the first convolution of the layer. If
-        greater than 1, this layer will downsample the input.
-      block_fn: Either `nn_blocks.ResidualBlock` or `nn_blocks.BottleneckBlock`.
-      block_repeats: `int` number of blocks contained in the layer.
-      stochastic_depth_drop_rate: `float` drop rate of the current block group.
-      name: `str`name for the block.
+      inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
+      filters: An `int` number of filters for the first convolution of the
+        layer.
+      strides: An `int` stride to use for the first convolution of the layer.
+        If greater than 1, this layer will downsample the input.
+      block_fn: The type of block group. Either `nn_blocks.ResidualBlock` or
+        `nn_blocks.BottleneckBlock`.
+      block_repeats: An `int` number of blocks contained in the layer.
+      stochastic_depth_drop_rate: A `float` of drop rate of the current block
+        group.
+      name: A `str` name for the block.
 
     Returns:
-      The output `Tensor` of the block layer.
+      The output `tf.Tensor` of the block layer.
     """
     x = block_fn(
         filters=filters,

official/vision/beta/modeling/backbones/resnet_3d.py

@@ -41,7 +41,7 @@ RESNET_SPECS = {
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class ResNet3D(tf.keras.Model):
-  """Class to build 3D ResNet family model."""
+  """Creates a 3D ResNet family model."""
 
   def __init__(self,
                model_id: int,
@@ -60,32 +60,33 @@ class ResNet3D(tf.keras.Model):
                kernel_regularizer=None,
                bias_regularizer=None,
                **kwargs):
-    """ResNet3D initialization function.
+    """Initializes a 3D ResNet model.
 
     Args:
-      model_id: `int` depth of ResNet backbone model.
-      temporal_strides: a list of integers that specifies the temporal strides
+      model_id: An `int` of depth of ResNet backbone model.
+      temporal_strides: A list of integers that specifies the temporal strides
         for all 3d blocks.
-      temporal_kernel_sizes: a list of tuples that specifies the temporal kernel
+      temporal_kernel_sizes: A list of tuples that specifies the temporal kernel
         sizes for all 3d blocks in different block groups.
-      use_self_gating: a list of booleans to specify applying self-gating module
+      use_self_gating: A list of booleans to specify applying self-gating module
         or not in each block group. If None, self-gating is not applied.
-      input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
-      stem_conv_temporal_kernel_size: `int` temporal kernel size for the first
-        conv layer.
-      stem_conv_temporal_stride: `int` temporal stride for the first conv layer.
-      stem_pool_temporal_stride: `int` temporal stride for the first pool layer.
-      activation: `str` name of the activation function.
-      use_sync_bn: if True, use synchronized batch normalization.
-      norm_momentum: `float` normalization omentum for the moving average.
-      norm_epsilon: `float` small float added to variance to avoid dividing by
-        zero.
-      kernel_initializer: kernel_initializer for convolutional layers.
-      kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
+      input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
+      stem_conv_temporal_kernel_size: An `int` of temporal kernel size for the
+        first conv layer.
+      stem_conv_temporal_stride: An `int` of temporal stride for the first conv
+        layer.
+      stem_pool_temporal_stride: An `int` of temporal stride for the first pool
+        layer.
+      activation: A `str` of name of the activation function.
+      use_sync_bn: If True, use synchronized batch normalization.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      kernel_initializer: A str for kernel initializer of convolutional layers.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default to None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
         Default to None.
-      bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
-        Default to None.
-      **kwargs: keyword arguments to be passed.
+      **kwargs: Additional keyword arguments to be passed.
     """
     self._model_id = model_id
     self._temporal_strides = temporal_strides
@@ -181,21 +182,23 @@ class ResNet3D(tf.keras.Model):
     """Creates one group of blocks for the ResNet3D model.
 
     Args:
-      inputs: `Tensor` of size `[batch, channels, height, width]`.
-      filters: `int` number of filters for the first convolution of the layer.
-      temporal_kernel_sizes: a tuple that specifies the temporal kernel sizes
+      inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
+      filters: An `int` of number of filters for the first convolution of the
+        layer.
+      temporal_kernel_sizes: A tuple that specifies the temporal kernel sizes
         for each block in the current group.
-      temporal_strides: `int` temporal strides for the first convolution in this
-        group.
-      spatial_strides: `int` stride to use for the first convolution of the
+      temporal_strides: An `int` of temporal strides for the first convolution
+        in this group.
+      spatial_strides: An `int` stride to use for the first convolution of the
         layer. If greater than 1, this layer will downsample the input.
       block_fn: Either `nn_blocks.ResidualBlock` or `nn_blocks.BottleneckBlock`.
-      block_repeats: `int` number of blocks contained in the layer.
-      use_self_gating: `bool` apply self-gating module or not.
-      name: `str`name for the block.
+      block_repeats: An `int` of number of blocks contained in the layer.
+      use_self_gating: A `bool` that specifies whether to apply self-gating
+        module or not.
+      name: A `str` name for the block.
 
     Returns:
-      The output `Tensor` of the block layer.
+      The output `tf.Tensor` of the block layer.
     """
     if len(temporal_kernel_sizes) != block_repeats:
       raise ValueError(

official/vision/beta/modeling/backbones/resnet_deeplab.py

@@ -45,12 +45,12 @@ RESNET_SPECS = {
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class DilatedResNet(tf.keras.Model):
-  """Class to build ResNet model with Deeplabv3 modifications.
+  """Creates a ResNet model with Deeplabv3 modifications.
 
-  This backbone is suitable for semantic segmentation. It was proposed in:
-  [1] Liang-Chieh Chen, George Papandreou, Florian Schroff, Hartwig Adam
+  This backbone is suitable for semantic segmentation. This implements
+    Liang-Chieh Chen, George Papandreou, Florian Schroff, Hartwig Adam.
     Rethinking Atrous Convolution for Semantic Image Segmentation.
-    arXiv:1706.05587
+    (https://arxiv.org/pdf/1706.05587)
   """
 
   def __init__(self,
@@ -70,30 +70,31 @@ class DilatedResNet(tf.keras.Model):
                kernel_regularizer=None,
                bias_regularizer=None,
                **kwargs):
-    """ResNet with DeepLab modification initialization function.
+    """Initializes a ResNet model with DeepLab modification.
 
     Args:
-      model_id: `int` depth of ResNet backbone model.
-      output_stride: `int` output stride, ratio of input to output resolution.
-      input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
-      stem_type: `standard` or `deeplab`, deeplab replaces 7x7 conv by 3 3x3
-        convs.
-      se_ratio: `float` or None. Ratio of the Squeeze-and-Excitation layer.
-      init_stochastic_depth_rate: `float` initial stochastic depth rate.
-      multigrid: `Tuple` of the same length as the number of blocks in the last
+      model_id: An `int` specifies depth of ResNet backbone model.
+      output_stride: An `int` of output stride, ratio of input to output
+        resolution.
+      input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
+      stem_type: A `str` of stem type. Can be `standard` or `deeplab`. `deeplab`
+        replaces 7x7 conv by 3 3x3 convs.
+      se_ratio: A `float` or None. Ratio of the Squeeze-and-Excitation layer.
+      init_stochastic_depth_rate: A `float` of initial stochastic depth rate.
+      multigrid: A tuple of the same length as the number of blocks in the last
         resnet stage.
-      last_stage_repeats: `int`, how many times last stage is repeated.
-      activation: `str` name of the activation function.
-      use_sync_bn: if True, use synchronized batch normalization.
-      norm_momentum: `float` normalization omentum for the moving average.
-      norm_epsilon: `float` small float added to variance to avoid dividing by
-        zero.
-      kernel_initializer: kernel_initializer for convolutional layers.
-      kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
-                          Default to None.
-      bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
-                        Default to None.
-      **kwargs: keyword arguments to be passed.
+      last_stage_repeats: An `int` that specifies how many times last stage is
+        repeated.
+      activation: A `str` name of the activation function.
+      use_sync_bn: If True, use synchronized batch normalization.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      kernel_initializer: A str for kernel initializer of convolutional layers.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default to None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
+        Default to None.
+      **kwargs: Additional keyword arguments to be passed.
     """
     self._model_id = model_id
     self._output_stride = output_stride
@@ -247,20 +248,22 @@ class DilatedResNet(tf.keras.Model):
     Deeplab applies strides at the last block.
 
     Args:
-      inputs: `Tensor` of size `[batch, channels, height, width]`.
-      filters: `int` number of filters for the first convolution of the layer.
-      strides: `int` stride to use for the first convolution of the layer. If
-        greater than 1, this layer will downsample the input.
-      dilation_rate: `int`, diluted convolution rates.
+      inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
+      filters: An `int` off number of filters for the first convolution of the
+        layer.
+      strides: An `int` of stride to use for the first convolution of the layer.
+        If greater than 1, this layer will downsample the input.
+      dilation_rate: An `int` of diluted convolution rates.
       block_fn: Either `nn_blocks.ResidualBlock` or `nn_blocks.BottleneckBlock`.
-      block_repeats: `int` number of blocks contained in the layer.
-      stochastic_depth_drop_rate: `float` drop rate of the current block group.
-      multigrid: List of ints or None, if specified, dilation rates for each
+      block_repeats: An `int` of number of blocks contained in the layer.
+      stochastic_depth_drop_rate: A `float` of drop rate of the current block
+        group.
+      multigrid: A list of `int` or None. If specified, dilation rates for each
         block is scaled up by its corresponding factor in the multigrid.
-      name: `str`name for the block.
+      name: A `str` name for the block.
 
     Returns:
-      The output `Tensor` of the block layer.
+      The output `tf.Tensor` of the block layer.
     """
     if multigrid is not None and len(multigrid) != block_repeats:
       raise ValueError('multigrid has to match number of block_repeats')

official/vision/beta/modeling/backbones/revnet.py

@@ -13,12 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =============================================================================="""
-"""RevNet Implementation.
-
-[1] Aidan N. Gomez, Mengye Ren, Raquel Urtasun, Roger B. Grosse
-    The Reversible Residual Network: Backpropagation Without Storing Activations
-    https://arxiv.org/pdf/1707.04585.pdf
-"""
+"""Contains definitions of RevNet."""
 
 from typing import Any, Callable, Dict, Optional
 # Import libraries
@@ -55,7 +50,14 @@ REVNET_SPECS = {
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class RevNet(tf.keras.Model):
-  """Reversible ResNet, RevNet implementation."""
+  """Creates a Reversible ResNet (RevNet) family model.
+
+  This implements:
+    Aidan N. Gomez, Mengye Ren, Raquel Urtasun, Roger B. Grosse.
+    The Reversible Residual Network: Backpropagation Without Storing
+    Activations.
+    (https://arxiv.org/pdf/1707.04585.pdf)
+  """
 
   def __init__(self,
                model_id: int,
@@ -68,19 +70,19 @@ class RevNet(tf.keras.Model):
                kernel_initializer: str = 'VarianceScaling',
                kernel_regularizer: tf.keras.regularizers.Regularizer = None,
                **kwargs):
-    """RevNet initialization function.
+    """Initializes a RevNet model.
 
     Args:
-      model_id: `int` depth/id of ResNet backbone model.
-      input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
-      activation: `str` name of the activation function.
-      use_sync_bn: `bool` if True, use synchronized batch normalization.
-      norm_momentum: `float` normalization omentum for the moving average.
-      norm_epsilon: `float` small float added to variance to avoid dividing by
-        zero.
-      kernel_initializer: `str` kernel_initializer for convolutional layers.
-      kernel_regularizer: `tf.keras.regularizers.Regularizer` for Conv2D.
-      **kwargs: additional keyword arguments to be passed.
+      model_id: An `int` of depth/id of ResNet backbone model.
+      input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
+      activation: A `str` name of the activation function.
+      use_sync_bn: If True, use synchronized batch normalization.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A `float` added to variance to avoid dividing by zero.
+      kernel_initializer: A str for kernel initializer of convolutional layers.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default to None.
+      **kwargs: Additional keyword arguments to be passed.
     """
     self._model_id = model_id
     self._input_specs = input_specs
@@ -148,19 +150,21 @@ class RevNet(tf.keras.Model):
     """Creates one reversible block for RevNet model.
 
     Args:
-      inputs: `Tensor` of size `[batch, channels, height, width]`.
-      filters: `int` number of filters for the first convolution of the layer.
-      strides: `int` stride to use for the first convolution of the layer. If
+      inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
+      filters: An `int` number of filters for the first convolution of the
+        layer.
+      strides: An `int` stride to use for the first convolution of the layer. If
         greater than 1, this block group will downsample the input.
       inner_block_fn: Either `nn_blocks.ResidualInner` or
         `nn_blocks.BottleneckResidualInner`.
-      block_repeats: `int` number of blocks contained in this block group.
-      batch_norm_first: `bool` whether to apply BatchNormalization and
-        activation layer before feeding into convolution layers.
-      name: `str`name for the block.
+      block_repeats: An `int` number of blocks contained in this block group.
+      batch_norm_first: A `bool` that specifies whether to apply
+        BatchNormalization and activation layer before feeding into convolution
+        layers.
+      name: A `str` name for the block.
 
     Returns:
-      The output `Tensor` of the block layer.
+      The output `tf.Tensor` of the block layer.
     """
     x = inputs
     for i in range(block_repeats):
@@ -210,7 +214,7 @@ def build_revnet(
     input_specs: tf.keras.layers.InputSpec,
     model_config,
     l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
-  """Builds ResNet 3d backbone from a config."""
+  """Builds RevNet backbone from a config."""
   backbone_type = model_config.backbone.type
   backbone_cfg = model_config.backbone.get()
   norm_activation_config = model_config.norm_activation

official/vision/beta/modeling/backbones/spinenet.py

@@ -13,12 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Implementation of SpineNet model.
-
-X. Du, T-Y. Lin, P. Jin, G. Ghiasi, M. Tan, Y. Cui, Q. V. Le, X. Song
-SpineNet: Learning Scale-Permuted Backbone for Recognition and Localization
-https://arxiv.org/abs/1912.05027
-"""
+"""Contains definitions of SpineNet Networks."""
 import math
 
 # Import libraries
@@ -117,7 +112,14 @@ def build_block_specs(block_specs=None):
 
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class SpineNet(tf.keras.Model):
-  """Class to build SpineNet models."""
+  """Creates a SpineNet family model.
+
+  This implements:
+    Xianzhi Du, Tsung-Yi Lin, Pengchong Jin, Golnaz Ghiasi, Mingxing Tan,
+    Yin Cui, Quoc V. Le, Xiaodan Song.
+    SpineNet: Learning Scale-Permuted Backbone for Recognition and Localization.
+    (https://arxiv.org/abs/1912.05027)
+  """
 
   def __init__(self,
                input_specs=tf.keras.layers.InputSpec(shape=[None, 640, 640, 3]),
@@ -137,7 +139,34 @@ class SpineNet(tf.keras.Model):
                norm_momentum=0.99,
                norm_epsilon=0.001,
                **kwargs):
-    """SpineNet model."""
+    """Initializes a SpineNet model.
+
+    Args:
+      input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
+      min_level: An `int` of min level for output mutiscale features.
+      max_level: An `int` of max level for output mutiscale features.
+      block_specs: The block specifications for the SpineNet model discovered by
+        NAS.
+      endpoints_num_filters: An `int` of feature dimension for the output
+        endpoints.
+      resample_alpha: A `float` of resampling factor in cross-scale connections.
+      block_repeats: An `int` of number of blocks contained in the layer.
+      filter_size_scale: A `float` of multiplier for the filters (number of
+        channels) for all convolution ops. The value must be greater than zero.
+        Typical usage will be to set this value in (0, 1) to reduce the number
+        of parameters or computation cost of the model.
+      init_stochastic_depth_rate: A `float` of initial stochastic depth rate.
+      kernel_initializer: A str for kernel initializer of convolutional layers.
+      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
+        Conv2D. Default to None.
+      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
+        Default to None.
+      activation: A `str` name of the activation function.
+      use_sync_bn: If True, use synchronized batch normalization.
+      norm_momentum: A `float` of normalization momentum for the moving average.
+      norm_epsilon: A small `float` added to variance to avoid dividing by zero.
+      **kwargs: Additional keyword arguments to be passed.
+    """
     self._input_specs = input_specs
     self._min_level = min_level
     self._max_level = max_level
@@ -235,7 +264,7 @@ class SpineNet(tf.keras.Model):
     return tf.identity(x, name=name)
 
   def _build_stem(self, inputs):
-    """Build SpineNet stem."""
+    """Builds SpineNet stem."""
     x = layers.Conv2D(
         filters=64,
         kernel_size=7,
@@ -271,7 +300,7 @@ class SpineNet(tf.keras.Model):
                                     net,
                                     input_width,
                                     weighted_fusion=False):
-    """Build scale-permuted network."""
+    """Builds scale-permuted network."""
     net_sizes = [int(math.ceil(input_width / 2**2))] * len(net)
     net_block_fns = [self._init_block_fn] * len(net)
     num_outgoing_connections = [0] * len(net)
@@ -363,7 +392,7 @@ class SpineNet(tf.keras.Model):
     return endpoints
 
   def _build_endpoints(self, net):
-    """Match filter size for endpoints before sharing conv layers."""
+    """Matches filter size for endpoints before sharing conv layers."""
     endpoints = {}
     for level in range(self._min_level, self._max_level + 1):
       x = layers.Conv2D(
@@ -392,7 +421,7 @@ class SpineNet(tf.keras.Model):
                            target_num_filters,
                            target_block_fn,
                            alpha=0.5):
-    """Match resolution and feature dimension."""
+    """Matches resolution and feature dimension."""
     _, _, _, input_num_filters = inputs.get_shape().as_list()
     if input_block_fn == 'bottleneck':
       input_num_filters /= 4
@@ -493,7 +522,7 @@ def build_spinenet(
     input_specs: tf.keras.layers.InputSpec,
     model_config,
     l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
-  """Builds ResNet 3d backbone from a config."""
+  """Builds SpineNet backbone from a config."""
   backbone_type = model_config.backbone.type
   backbone_cfg = model_config.backbone.get()
   norm_activation_config = model_config.norm_activation