Add more pytype checking.

PiperOrigin-RevId: 368129317

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

@@ -15,8 +15,13 @@
 """Contains definitions of EfficientNet Networks."""
 
 import math
+from typing import Any, List, Tuple
+
 # Import libraries
+
 import tensorflow as tf
+
+from official.modeling import hyperparams
 from official.modeling import tf_utils
 from official.vision.beta.modeling.backbones import factory
 from official.vision.beta.modeling.layers import nn_blocks
@@ -50,14 +55,32 @@ SCALING_MAP = {
 }
 
 
-def round_repeats(repeats, multiplier, skip=False):
+class BlockSpec():
+  """A container class that specifies the block configuration for MnasNet."""
+
+  def __init__(self, block_fn: str, block_repeats: int, kernel_size: int,
+               strides: int, expand_ratio: float, in_filters: int,
+               out_filters: int, is_output: bool, width_scale: float,
+               depth_scale: float):
+    self.block_fn = block_fn
+    self.block_repeats = round_repeats(block_repeats, depth_scale)
+    self.kernel_size = kernel_size
+    self.strides = strides
+    self.expand_ratio = expand_ratio
+    self.in_filters = nn_layers.round_filters(in_filters, width_scale)
+    self.out_filters = nn_layers.round_filters(out_filters, width_scale)
+    self.is_output = is_output
+
+
+def round_repeats(repeats: int, multiplier: float, skip: bool = False) -> int:
   """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):
+def block_spec_decoder(specs: List[Tuple[Any, ...]], width_scale: float,
+                       depth_scale: float) -> List[BlockSpec]:
   """Decodes and returns specs for a block."""
   decoded_specs = []
   for s in specs:
@@ -69,22 +92,6 @@ def block_spec_decoder(specs, width_scale, depth_scale):
   return decoded_specs
 
 
-class BlockSpec(object):
-  """A container class that specifies the block configuration for MnasNet."""
-
-  def __init__(self, block_fn, block_repeats, kernel_size, strides,
-               expand_ratio, in_filters, out_filters, is_output, width_scale,
-               depth_scale):
-    self.block_fn = block_fn
-    self.block_repeats = round_repeats(block_repeats, depth_scale)
-    self.kernel_size = kernel_size
-    self.strides = strides
-    self.expand_ratio = expand_ratio
-    self.in_filters = nn_layers.round_filters(in_filters, width_scale)
-    self.out_filters = nn_layers.round_filters(out_filters, width_scale)
-    self.is_output = is_output
-
-
 @tf.keras.utils.register_keras_serializable(package='Vision')
 class EfficientNet(tf.keras.Model):
   """Creates an EfficientNet family model.
@@ -96,17 +103,18 @@ class EfficientNet(tf.keras.Model):
   """
 
   def __init__(self,
-               model_id,
-               input_specs=layers.InputSpec(shape=[None, None, None, 3]),
-               se_ratio=0.0,
-               stochastic_depth_drop_rate=0.0,
-               kernel_initializer='VarianceScaling',
-               kernel_regularizer=None,
-               bias_regularizer=None,
-               activation='relu',
-               use_sync_bn=False,
-               norm_momentum=0.99,
-               norm_epsilon=0.001,
+               model_id: str,
+               input_specs: tf.keras.layers.InputSpec = layers.InputSpec(
+                   shape=[None, None, None, 3]),
+               se_ratio: float = 0.0,
+               stochastic_depth_drop_rate: float = 0.0,
+               kernel_initializer: str = 'VarianceScaling',
+               kernel_regularizer: tf.keras.regularizers.Regularizer = None,
+               bias_regularizer: tf.keras.regularizers.Regularizer = None,
+               activation: str = 'relu',
+               use_sync_bn: bool = False,
+               norm_momentum: float = 0.99,
+               norm_epsilon: float = 0.001,
                **kwargs):
     """Initializes an EfficientNet model.
 
@@ -205,7 +213,10 @@ class EfficientNet(tf.keras.Model):
     super(EfficientNet, self).__init__(
         inputs=inputs, outputs=endpoints, **kwargs)
 
-  def _block_group(self, inputs, specs, name='block_group'):
+  def _block_group(self,
+                   inputs: tf.Tensor,
+                   specs: BlockSpec,
+                   name: str = 'block_group'):
     """Creates one group of blocks for the EfficientNet model.
 
     Args:
@@ -286,7 +297,7 @@ class EfficientNet(tf.keras.Model):
 @factory.register_backbone_builder('efficientnet')
 def build_efficientnet(
     input_specs: tf.keras.layers.InputSpec,
-    model_config,
+    model_config: hyperparams.Config,
     l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
   """Builds EfficientNet backbone from a config."""
   backbone_type = model_config.backbone.type

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

@@ -43,9 +43,11 @@ in place that uses it.
 
 """
 # Import libraries
+
 import tensorflow as tf
 
 from official.core import registry
+from official.modeling import hyperparams
 
 
 _REGISTERED_BACKBONE_CLS = {}
@@ -79,9 +81,10 @@ def register_backbone_builder(key: str):
   return registry.register(_REGISTERED_BACKBONE_CLS, key)
 
 
-def build_backbone(input_specs: tf.keras.layers.InputSpec,
-                   model_config,
-                   l2_regularizer: tf.keras.regularizers.Regularizer = None):
+def build_backbone(
+    input_specs: tf.keras.layers.InputSpec,
+    model_config: hyperparams.Config,
+    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
   """Builds backbone from a config.
 
   Args:

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

@@ -26,7 +26,6 @@ from official.vision.beta.modeling.layers import nn_blocks
 from official.vision.beta.modeling.layers import nn_layers
 
 layers = tf.keras.layers
-regularizers = tf.keras.regularizers
 
 
 #  pylint: disable=pointless-string-statement
@@ -417,18 +416,19 @@ class BlockSpec(hyperparams.Config):
   use_bias: bool = False
   use_normalization: bool = True
   activation: str = 'relu6'
-  # used for block type InvertedResConv
+  # Used for block type InvertedResConv.
   expand_ratio: Optional[float] = 6.
-  # used for block type InvertedResConv with SE
+  # Used for block type InvertedResConv with SE.
   se_ratio: Optional[float] = None
   use_depthwise: bool = True
   use_residual: bool = True
   is_output: bool = True
 
 
-def block_spec_decoder(specs: Dict[Any, Any],
+def block_spec_decoder(
+    specs: Dict[Any, Any],
     filter_size_scale: float,
-                       # set to 1 for mobilenetv1
+    # Set to 1 for mobilenetv1.
     divisible_by: int = 8,
     finegrain_classification_mode: bool = True):
   """Decodes specs for a block.
@@ -491,23 +491,23 @@ class MobileNet(tf.keras.Model):
       self,
       model_id: str = 'MobileNetV2',
       filter_size_scale: float = 1.0,
-      input_specs: layers.InputSpec = layers.InputSpec(
+      input_specs: tf.keras.layers.InputSpec = layers.InputSpec(
           shape=[None, None, None, 3]),
-      # The followings are for hyper-parameter tuning
+      # 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
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf.keras.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 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 is not used in MobileNetV1.
       finegrain_classification_mode: bool = True,
       **kwargs):
     """Initializes a MobileNet model.
@@ -636,8 +636,8 @@ class MobileNet(tf.keras.Model):
       # A small catch for gpooling block with None strides
       if not block_def.strides:
         block_def.strides = 1
-      if self._output_stride is not None \
-          and current_stride == self._output_stride:
+      if (self._output_stride is not None and
+          current_stride == self._output_stride):
         # If we have reached the target output_stride, then we need to employ
         # atrous convolution with stride=1 and multiply the atrous rate by the
         # current unit's stride for use in subsequent layers.
@@ -764,7 +764,7 @@ class MobileNet(tf.keras.Model):
 @factory.register_backbone_builder('mobilenet')
 def build_mobilenet(
     input_specs: tf.keras.layers.InputSpec,
-    model_config,
+    model_config: hyperparams.Config,
     l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
   """Builds MobileNet backbone from a config."""
   backbone_type = model_config.backbone.type

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

@@ -14,8 +14,12 @@
 
 """Contains definitions of ResNet and ResNet-RS models."""
 
+from typing import Callable, Optional
+
 # Import libraries
 import tensorflow as tf
+
+from official.modeling import hyperparams
 from official.modeling import tf_utils
 from official.vision.beta.modeling.backbones import factory
 from official.vision.beta.modeling.layers import nn_blocks
@@ -99,22 +103,24 @@ class ResNet(tf.keras.Model):
     (https://arxiv.org/abs/2103.07579).
   """
 
-  def __init__(self,
-               model_id,
-               input_specs=layers.InputSpec(shape=[None, None, None, 3]),
-               depth_multiplier=1.0,
-               stem_type='v0',
-               resnetd_shortcut=False,
-               replace_stem_max_pool=False,
-               se_ratio=None,
-               init_stochastic_depth_rate=0.0,
-               activation='relu',
-               use_sync_bn=False,
-               norm_momentum=0.99,
-               norm_epsilon=0.001,
-               kernel_initializer='VarianceScaling',
-               kernel_regularizer=None,
-               bias_regularizer=None,
+  def __init__(
+      self,
+      model_id: int,
+      input_specs: tf.keras.layers.InputSpec = layers.InputSpec(
+          shape=[None, None, None, 3]),
+      depth_multiplier: float = 1.0,
+      stem_type: str = 'v0',
+      resnetd_shortcut: bool = False,
+      replace_stem_max_pool: bool = False,
+      se_ratio: Optional[float] = None,
+      init_stochastic_depth_rate: float = 0.0,
+      activation: str = 'relu',
+      use_sync_bn: bool = False,
+      norm_momentum: float = 0.99,
+      norm_epsilon: float = 0.001,
+      kernel_initializer: str = 'VarianceScaling',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
       **kwargs):
     """Initializes a ResNet model.
 
@@ -274,13 +280,13 @@ class ResNet(tf.keras.Model):
     super(ResNet, self).__init__(inputs=inputs, outputs=endpoints, **kwargs)
 
   def _block_group(self,
-                   inputs,
-                   filters,
-                   strides,
-                   block_fn,
-                   block_repeats=1,
-                   stochastic_depth_drop_rate=0.0,
-                   name='block_group'):
+                   inputs: tf.Tensor,
+                   filters: int,
+                   strides: int,
+                   block_fn: Callable[..., tf.keras.layers.Layer],
+                   block_repeats: int = 1,
+                   stochastic_depth_drop_rate: float = 0.0,
+                   name: str = 'block_group'):
     """Creates one group of blocks for the ResNet model.
 
     Args:
@@ -366,7 +372,7 @@ class ResNet(tf.keras.Model):
 @factory.register_backbone_builder('resnet')
 def build_resnet(
     input_specs: tf.keras.layers.InputSpec,
-    model_config,
+    model_config: hyperparams.Config,
     l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
   """Builds ResNet backbone from a config."""
   backbone_type = model_config.backbone.type

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

@@ -13,10 +13,12 @@
 # limitations under the License.
 
 """Contains definitions of 3D Residual Networks."""
-from typing import List, Tuple
+from typing import Callable, List, Tuple, Optional
 
 # Import libraries
 import tensorflow as tf
+
+from official.modeling import hyperparams
 from official.modeling import tf_utils
 from official.vision.beta.modeling.backbones import factory
 from official.vision.beta.modeling.layers import nn_blocks_3d
@@ -74,25 +76,27 @@ RESNET_SPECS = {
 class ResNet3D(tf.keras.Model):
   """Creates a 3D ResNet family model."""
 
-  def __init__(self,
+  def __init__(
+      self,
       model_id: int,
       temporal_strides: List[int],
       temporal_kernel_sizes: List[Tuple[int]],
       use_self_gating: List[int] = None,
-               input_specs=layers.InputSpec(shape=[None, None, None, None, 3]),
-               stem_type='v0',
-               stem_conv_temporal_kernel_size=5,
-               stem_conv_temporal_stride=2,
-               stem_pool_temporal_stride=2,
-               init_stochastic_depth_rate=0.0,
-               activation='relu',
-               se_ratio=None,
-               use_sync_bn=False,
-               norm_momentum=0.99,
-               norm_epsilon=0.001,
-               kernel_initializer='VarianceScaling',
-               kernel_regularizer=None,
-               bias_regularizer=None,
+      input_specs: tf.keras.layers.InputSpec = layers.InputSpec(
+          shape=[None, None, None, None, 3]),
+      stem_type: str = 'v0',
+      stem_conv_temporal_kernel_size: int = 5,
+      stem_conv_temporal_stride: int = 2,
+      stem_pool_temporal_stride: int = 2,
+      init_stochastic_depth_rate: float = 0.0,
+      activation: str = 'relu',
+      se_ratio: Optional[float] = None,
+      use_sync_bn: bool = False,
+      norm_momentum: float = 0.99,
+      norm_epsilon: float = 0.001,
+      kernel_initializer: str = 'VarianceScaling',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
       **kwargs):
     """Initializes a 3D ResNet model.
 
@@ -259,16 +263,18 @@ class ResNet3D(tf.keras.Model):
     super(ResNet3D, self).__init__(inputs=inputs, outputs=endpoints, **kwargs)
 
   def _block_group(self,
-                   inputs,
-                   filters,
-                   temporal_kernel_sizes,
-                   temporal_strides,
-                   spatial_strides,
-                   block_fn=nn_blocks_3d.BottleneckBlock3D,
-                   block_repeats=1,
-                   stochastic_depth_drop_rate=0.0,
-                   use_self_gating=False,
-                   name='block_group'):
+                   inputs: tf.Tensor,
+                   filters: int,
+                   temporal_kernel_sizes: Tuple[int],
+                   temporal_strides: int,
+                   spatial_strides: int,
+                   block_fn: Callable[
+                       ...,
+                       tf.keras.layers.Layer] = nn_blocks_3d.BottleneckBlock3D,
+                   block_repeats: int = 1,
+                   stochastic_depth_drop_rate: float = 0.0,
+                   use_self_gating: bool = False,
+                   name: str = 'block_group'):
     """Creates one group of blocks for the ResNet3D model.
 
     Args:
@@ -410,7 +416,7 @@ def build_resnet3d(
 @factory.register_backbone_builder('resnet_3d_rs')
 def build_resnet3d_rs(
     input_specs: tf.keras.layers.InputSpec,
-    model_config,
+    model_config: hyperparams.Config,
     l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
   """Builds ResNet-3D-RS backbone from a config."""
   backbone_cfg = model_config.backbone.get()

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

@@ -14,6 +14,8 @@
 
 """Contains definitions of Residual Networks with Deeplab modifications."""
 
+from typing import Callable, Optional, Tuple, List
+
 import numpy as np
 import tensorflow as tf
 from official.modeling import tf_utils
@@ -53,22 +55,24 @@ class DilatedResNet(tf.keras.Model):
     (https://arxiv.org/pdf/1706.05587)
   """
 
-  def __init__(self,
-               model_id,
-               output_stride,
-               input_specs=layers.InputSpec(shape=[None, None, None, 3]),
-               stem_type='v0',
-               se_ratio=None,
-               init_stochastic_depth_rate=0.0,
-               multigrid=None,
-               last_stage_repeats=1,
-               activation='relu',
-               use_sync_bn=False,
-               norm_momentum=0.99,
-               norm_epsilon=0.001,
-               kernel_initializer='VarianceScaling',
-               kernel_regularizer=None,
-               bias_regularizer=None,
+  def __init__(
+      self,
+      model_id: int,
+      output_stride: int,
+      input_specs: tf.keras.layers.InputSpec = layers.InputSpec(
+          shape=[None, None, None, 3]),
+      stem_type: str = 'v0',
+      se_ratio: Optional[float] = None,
+      init_stochastic_depth_rate: float = 0.0,
+      multigrid: Optional[Tuple[int]] = None,
+      last_stage_repeats: int = 1,
+      activation: str = 'relu',
+      use_sync_bn: bool = False,
+      norm_momentum: float = 0.99,
+      norm_epsilon: float = 0.001,
+      kernel_initializer: str = 'VarianceScaling',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
       **kwargs):
     """Initializes a ResNet model with DeepLab modification.
 
@@ -234,15 +238,15 @@ class DilatedResNet(tf.keras.Model):
         inputs=inputs, outputs=endpoints, **kwargs)
 
   def _block_group(self,
-                   inputs,
-                   filters,
-                   strides,
-                   dilation_rate,
-                   block_fn,
-                   block_repeats=1,
-                   stochastic_depth_drop_rate=0.0,
-                   multigrid=None,
-                   name='block_group'):
+                   inputs: tf.Tensor,
+                   filters: int,
+                   strides: int,
+                   dilation_rate: int,
+                   block_fn: Callable[..., tf.keras.layers.Layer],
+                   block_repeats: int = 1,
+                   stochastic_depth_drop_rate: float = 0.0,
+                   multigrid: Optional[List[int]] = None,
+                   name: str = 'block_group'):
     """Creates one group of blocks for the ResNet model.
 
     Deeplab applies strides at the last block.

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

@@ -59,16 +59,17 @@ class RevNet(tf.keras.Model):
     (https://arxiv.org/pdf/1707.04585.pdf)
   """
 
-  def __init__(self,
+  def __init__(
+      self,
       model_id: int,
-               input_specs: tf.keras.layers.InputSpec
-               = tf.keras.layers.InputSpec(shape=[None, None, None, 3]),
+      input_specs: tf.keras.layers.InputSpec = tf.keras.layers.InputSpec(
+          shape=[None, None, None, 3]),
       activation: str = 'relu',
       use_sync_bn: bool = False,
       norm_momentum: float = 0.99,
       norm_epsilon: float = 0.001,
       kernel_initializer: str = 'VarianceScaling',
-               kernel_regularizer: tf.keras.regularizers.Regularizer = None,
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
       **kwargs):
     """Initializes a RevNet model.
 

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

@@ -12,13 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Lint as: python3
 """Contains definitions of SpineNet Networks."""
+
 import math
+from typing import Any, List, Optional, Tuple
 
 # Import libraries
+
 from absl import logging
 import tensorflow as tf
+
 from official.modeling import tf_utils
 from official.vision.beta.modeling.backbones import factory
 from official.vision.beta.modeling.layers import nn_blocks
@@ -95,14 +98,16 @@ SCALING_MAP = {
 class BlockSpec(object):
   """A container class that specifies the block configuration for SpineNet."""
 
-  def __init__(self, level, block_fn, input_offsets, is_output):
+  def __init__(self, level: int, block_fn: str, input_offsets: Tuple[int, int],
+               is_output: bool):
     self.level = level
     self.block_fn = block_fn
     self.input_offsets = input_offsets
     self.is_output = is_output
 
 
-def build_block_specs(block_specs=None):
+def build_block_specs(
+    block_specs: Optional[List[Tuple[Any, ...]]] = None) -> List[BlockSpec]:
   """Builds the list of BlockSpec objects for SpineNet."""
   if not block_specs:
     block_specs = SPINENET_BLOCK_SPECS
@@ -121,23 +126,25 @@ class SpineNet(tf.keras.Model):
     (https://arxiv.org/abs/1912.05027)
   """
 
-  def __init__(self,
-               input_specs=tf.keras.layers.InputSpec(shape=[None, 640, 640, 3]),
-               min_level=3,
-               max_level=7,
-               block_specs=build_block_specs(),
-               endpoints_num_filters=256,
-               resample_alpha=0.5,
-               block_repeats=1,
-               filter_size_scale=1.0,
-               init_stochastic_depth_rate=0.0,
-               kernel_initializer='VarianceScaling',
-               kernel_regularizer=None,
-               bias_regularizer=None,
-               activation='relu',
-               use_sync_bn=False,
-               norm_momentum=0.99,
-               norm_epsilon=0.001,
+  def __init__(
+      self,
+      input_specs: tf.keras.layers.InputSpec = tf.keras.layers.InputSpec(
+          shape=[None, 640, 640, 3]),
+      min_level: int = 3,
+      max_level: int = 7,
+      block_specs: List[BlockSpec] = build_block_specs(),
+      endpoints_num_filters: int = 256,
+      resample_alpha: float = 0.5,
+      block_repeats: int = 1,
+      filter_size_scale: float = 1.0,
+      init_stochastic_depth_rate: float = 0.0,
+      kernel_initializer: str = 'VarianceScaling',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      activation: str = 'relu',
+      use_sync_bn: bool = False,
+      norm_momentum: float = 0.99,
+      norm_epsilon: float = 0.001,
       **kwargs):
     """Initializes a SpineNet model.
 
@@ -145,8 +152,8 @@ class SpineNet(tf.keras.Model):
       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.
+      block_specs: A list of 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.
@@ -214,13 +221,13 @@ class SpineNet(tf.keras.Model):
     super(SpineNet, self).__init__(inputs=inputs, outputs=endpoints)
 
   def _block_group(self,
-                   inputs,
-                   filters,
-                   strides,
-                   block_fn_cand,
-                   block_repeats=1,
-                   stochastic_depth_drop_rate=None,
-                   name='block_group'):
+                   inputs: tf.Tensor,
+                   filters: int,
+                   strides: int,
+                   block_fn_cand: str,
+                   block_repeats: int = 1,
+                   stochastic_depth_drop_rate: Optional[float] = None,
+                   name: str = 'block_group'):
     """Creates one group of blocks for the SpineNet model."""
     block_fn_candidates = {
         'bottleneck': nn_blocks.BottleneckBlock,

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

@@ -29,10 +29,13 @@
 # ==============================================================================
 """Contains definitions of Mobile SpineNet Networks."""
 import math
+from typing import Any, List, Optional, Tuple
 
 # Import libraries
+
 from absl import logging
 import tensorflow as tf
+
 from official.modeling import tf_utils
 from official.vision.beta.modeling.backbones import factory
 from official.vision.beta.modeling.layers import nn_blocks
@@ -96,14 +99,16 @@ SCALING_MAP = {
 class BlockSpec(object):
   """A container class that specifies the block configuration for SpineNet."""
 
-  def __init__(self, level, block_fn, input_offsets, is_output):
+  def __init__(self, level: int, block_fn: str, input_offsets: Tuple[int, int],
+               is_output: bool):
     self.level = level
     self.block_fn = block_fn
     self.input_offsets = input_offsets
     self.is_output = is_output
 
 
-def build_block_specs(block_specs=None):
+def build_block_specs(
+    block_specs: Optional[List[Tuple[Any, ...]]] = None) -> List[BlockSpec]:
   """Builds the list of BlockSpec objects for SpineNet."""
   if not block_specs:
     block_specs = SPINENET_BLOCK_SPECS
@@ -126,24 +131,26 @@ class SpineNetMobile(tf.keras.Model):
     (https://arxiv.org/abs/2010.11426).
   """
 
-  def __init__(self,
-               input_specs=tf.keras.layers.InputSpec(shape=[None, 512, 512, 3]),
-               min_level=3,
-               max_level=7,
-               block_specs=build_block_specs(),
-               endpoints_num_filters=48,
-               se_ratio=0.2,
-               block_repeats=1,
-               filter_size_scale=1.0,
-               expand_ratio=6,
+  def __init__(
+      self,
+      input_specs: tf.keras.layers.InputSpec = tf.keras.layers.InputSpec(
+          shape=[None, 512, 512, 3]),
+      min_level: int = 3,
+      max_level: int = 7,
+      block_specs: List[BlockSpec] = build_block_specs(),
+      endpoints_num_filters: int = 256,
+      se_ratio: float = 0.2,
+      block_repeats: int = 1,
+      filter_size_scale: float = 1.0,
+      expand_ratio: int = 6,
       init_stochastic_depth_rate=0.0,
-               kernel_initializer='VarianceScaling',
-               kernel_regularizer=None,
-               bias_regularizer=None,
-               activation='relu',
-               use_sync_bn=False,
-               norm_momentum=0.99,
-               norm_epsilon=0.001,
+      kernel_initializer: str = 'VarianceScaling',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      activation: str = 'relu',
+      use_sync_bn: bool = False,
+      norm_momentum: float = 0.99,
+      norm_epsilon: float = 0.001,
       **kwargs):
     """Initializes a Mobile SpineNet model.
 
@@ -222,15 +229,15 @@ class SpineNetMobile(tf.keras.Model):
     super().__init__(inputs=inputs, outputs=endpoints)
 
   def _block_group(self,
-                   inputs,
-                   in_filters,
-                   out_filters,
-                   strides,
-                   expand_ratio=6,
-                   block_repeats=1,
-                   se_ratio=0.2,
-                   stochastic_depth_drop_rate=None,
-                   name='block_group'):
+                   inputs: tf.Tensor,
+                   in_filters: int,
+                   out_filters: int,
+                   strides: int,
+                   expand_ratio: int = 6,
+                   block_repeats: int = 1,
+                   se_ratio: float = 0.2,
+                   stochastic_depth_drop_rate: Optional[float] = None,
+                   name: str = 'block_group'):
     """Creates one group of blocks for the SpineNet model."""
     x = nn_blocks.InvertedBottleneckBlock(
         in_filters=in_filters,

official/vision/beta/modeling/classification_model.py

@@ -14,6 +14,7 @@
 
 """Build classification models."""
 
+from typing import Any, Mapping, Optional
 # Import libraries
 import tensorflow as tf
 
@@ -24,15 +25,17 @@ layers = tf.keras.layers
 class ClassificationModel(tf.keras.Model):
   """A classification class builder."""
 
-  def __init__(self,
-               backbone,
-               num_classes,
-               input_specs=layers.InputSpec(shape=[None, None, None, 3]),
-               dropout_rate=0.0,
-               kernel_initializer='random_uniform',
-               kernel_regularizer=None,
-               bias_regularizer=None,
-               add_head_batch_norm=False,
+  def __init__(
+      self,
+      backbone: tf.keras.Model,
+      num_classes: int,
+      input_specs: tf.keras.layers.InputSpec = layers.InputSpec(
+          shape=[None, None, None, 3]),
+      dropout_rate: float = 0.0,
+      kernel_initializer: str = 'random_uniform',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      add_head_batch_norm: bool = False,
       use_sync_bn: bool = False,
       norm_momentum: float = 0.99,
       norm_epsilon: float = 0.001,
@@ -103,15 +106,15 @@ class ClassificationModel(tf.keras.Model):
     self._norm = norm
 
   @property
-  def checkpoint_items(self):
+  def checkpoint_items(self) -> Mapping[str, tf.keras.Model]:
     """Returns a dictionary of items to be additionally checkpointed."""
     return dict(backbone=self.backbone)
 
   @property
-  def backbone(self):
+  def backbone(self) -> tf.keras.Model:
     return self._backbone
 
-  def get_config(self):
+  def get_config(self) -> Mapping[str, Any]:
     return self._config_dict
 
   @classmethod

official/vision/beta/modeling/decoders/aspp.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Contains definitions of Atrous Spatial Pyramid Pooling (ASPP) decoder."""
+from typing import Any, List, Optional, Mapping
 
 # Import libraries
 import tensorflow as tf
@@ -24,19 +25,20 @@ from official.vision import keras_cv
 class ASPP(tf.keras.layers.Layer):
   """Creates an Atrous Spatial Pyramid Pooling (ASPP) layer."""
 
-  def __init__(self,
-               level,
-               dilation_rates,
-               num_filters=256,
-               pool_kernel_size=None,
-               use_sync_bn=False,
-               norm_momentum=0.99,
-               norm_epsilon=0.001,
-               activation='relu',
-               dropout_rate=0.0,
-               kernel_initializer='VarianceScaling',
-               kernel_regularizer=None,
-               interpolation='bilinear',
+  def __init__(
+      self,
+      level: int,
+      dilation_rates: List[int],
+      num_filters: int = 256,
+      pool_kernel_size: Optional[int] = None,
+      use_sync_bn: bool = False,
+      norm_momentum: float = 0.99,
+      norm_epsilon: float = 0.001,
+      activation: str = 'relu',
+      dropout_rate: float = 0.0,
+      kernel_initializer: str = 'VarianceScaling',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      interpolation: str = 'bilinear',
       **kwargs):
     """Initializes an Atrous Spatial Pyramid Pooling (ASPP) layer.
 
@@ -97,7 +99,7 @@ class ASPP(tf.keras.layers.Layer):
         kernel_regularizer=self._config_dict['kernel_regularizer'],
         interpolation=self._config_dict['interpolation'])
 
-  def call(self, inputs):
+  def call(self, inputs: Mapping[str, tf.Tensor]) -> Mapping[str, tf.Tensor]:
     """Calls the Atrous Spatial Pyramid Pooling (ASPP) layer on an input.
 
     The output of ASPP will be a dict of {`level`, `tf.Tensor`} even if only one
@@ -120,7 +122,7 @@ class ASPP(tf.keras.layers.Layer):
     outputs[level] = self.aspp(inputs[level])
     return outputs
 
-  def get_config(self):
+  def get_config(self) -> Mapping[str, Any]:
     return self._config_dict
 
   @classmethod

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

@@ -15,15 +15,21 @@
 # Lint as: python3
 """Contains the factory method to create decoders."""
 
+from typing import Mapping, Optional
+
 # Import libraries
+
 import tensorflow as tf
 
+from official.modeling import hyperparams
 from official.vision.beta.modeling import decoders
 
 
-def build_decoder(input_specs,
-                  model_config,
-                  l2_regularizer: tf.keras.regularizers.Regularizer = None):
+def build_decoder(
+    input_specs: Mapping[str, tf.TensorShape],
+    model_config: hyperparams.Config,
+    l2_regularizer: Optional[tf.keras.regularizers.Regularizer] = None
+) -> tf.keras.Model:
   """Builds decoder from a config.
 
   Args:

official/vision/beta/modeling/decoders/fpn.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Contains the definitions of Feature Pyramid Networks (FPN)."""
+from typing import Any, Mapping, Optional
 
 # Import libraries
 import tensorflow as tf
@@ -32,19 +33,20 @@ class FPN(tf.keras.Model):
   (https://arxiv.org/pdf/1612.03144)
   """
 
-  def __init__(self,
-               input_specs,
-               min_level=3,
-               max_level=7,
-               num_filters=256,
-               use_separable_conv=False,
-               activation='relu',
-               use_sync_bn=False,
-               norm_momentum=0.99,
-               norm_epsilon=0.001,
-               kernel_initializer='VarianceScaling',
-               kernel_regularizer=None,
-               bias_regularizer=None,
+  def __init__(
+      self,
+      input_specs: Mapping[str, tf.TensorShape],
+      min_level: int = 3,
+      max_level: int = 7,
+      num_filters: int = 256,
+      use_separable_conv: bool = False,
+      activation: str = 'relu',
+      use_sync_bn: bool = False,
+      norm_momentum: float = 0.99,
+      norm_epsilon: float = 0.001,
+      kernel_initializer: str = 'VarianceScaling',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
       **kwargs):
     """Initializes a Feature Pyramid Network (FPN).
 
@@ -162,7 +164,8 @@ class FPN(tf.keras.Model):
 
     super(FPN, self).__init__(inputs=inputs, outputs=feats, **kwargs)
 
-  def _build_input_pyramid(self, input_specs, min_level):
+  def _build_input_pyramid(self, input_specs: Mapping[str, tf.TensorShape],
+                           min_level: int):
     assert isinstance(input_specs, dict)
     if min(input_specs.keys()) > str(min_level):
       raise ValueError(
@@ -173,7 +176,7 @@ class FPN(tf.keras.Model):
       inputs[level] = tf.keras.Input(shape=spec[1:])
     return inputs
 
-  def get_config(self):
+  def get_config(self) -> Mapping[str, Any]:
     return self._config_dict
 
   @classmethod
@@ -181,6 +184,6 @@ class FPN(tf.keras.Model):
     return cls(**config)
 
   @property
-  def output_specs(self):
+  def output_specs(self) -> Mapping[str, tf.TensorShape]:
     """A dict of {level: TensorShape} pairs for the model output."""
     return self._output_specs

official/vision/beta/modeling/decoders/nasfpn.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Contains definitions of NAS-FPN."""
+from typing import Any, Mapping, List, Tuple, Optional
 
 # Import libraries
 from absl import logging
@@ -35,17 +36,19 @@ NASFPN_BLOCK_SPECS = [
 ]
 
 
-class BlockSpec(object):
+class BlockSpec():
   """A container class that specifies the block configuration for NAS-FPN."""
 
-  def __init__(self, level, combine_fn, input_offsets, is_output):
+  def __init__(self, level: int, combine_fn: str,
+               input_offsets: Tuple[int, int], is_output: bool):
     self.level = level
     self.combine_fn = combine_fn
     self.input_offsets = input_offsets
     self.is_output = is_output
 
 
-def build_block_specs(block_specs=None):
+def build_block_specs(
+    block_specs: Optional[List[Tuple[Any, ...]]] = None) -> List[BlockSpec]:
   """Builds the list of BlockSpec objects for NAS-FPN."""
   if not block_specs:
     block_specs = NASFPN_BLOCK_SPECS
@@ -63,21 +66,22 @@ class NASFPN(tf.keras.Model):
   (https://arxiv.org/abs/1904.07392)
   """
 
-  def __init__(self,
-               input_specs,
-               min_level=3,
-               max_level=7,
-               block_specs=build_block_specs(),
-               num_filters=256,
-               num_repeats=5,
-               use_separable_conv=False,
-               activation='relu',
-               use_sync_bn=False,
-               norm_momentum=0.99,
-               norm_epsilon=0.001,
-               kernel_initializer='VarianceScaling',
-               kernel_regularizer=None,
-               bias_regularizer=None,
+  def __init__(
+      self,
+      input_specs: Mapping[str, tf.TensorShape],
+      min_level: int = 3,
+      max_level: int = 7,
+      block_specs: List[BlockSpec] = build_block_specs(),
+      num_filters: int = 256,
+      num_repeats: int = 5,
+      use_separable_conv: bool = False,
+      activation: str = 'relu',
+      use_sync_bn: bool = False,
+      norm_momentum: float = 0.99,
+      norm_epsilon: float = 0.001,
+      kernel_initializer: str = 'VarianceScaling',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
       **kwargs):
     """Initializes a NAS-FPN model.
 
@@ -191,7 +195,8 @@ class NASFPN(tf.keras.Model):
                     for level in output_feats.keys()}
     super(NASFPN, self).__init__(inputs=inputs, outputs=output_feats, **kwargs)
 
-  def _build_input_pyramid(self, input_specs, min_level):
+  def _build_input_pyramid(self, input_specs: Mapping[str, tf.TensorShape],
+                           min_level: int):
     assert isinstance(input_specs, dict)
     if min(input_specs.keys()) > str(min_level):
       raise ValueError(
@@ -300,7 +305,7 @@ class NASFPN(tf.keras.Model):
     logging.info('Output feature pyramid: %s', output_feats)
     return output_feats
 
-  def get_config(self):
+  def get_config(self) -> Mapping[str, Any]:
     return self._config_dict
 
   @classmethod
@@ -308,6 +313,6 @@ class NASFPN(tf.keras.Model):
     return cls(**config)
 
   @property
-  def output_specs(self):
+  def output_specs(self) -> Mapping[str, tf.TensorShape]:
     """A dict of {level: TensorShape} pairs for the model output."""
     return self._output_specs

official/vision/beta/modeling/factory.py

@@ -42,7 +42,7 @@ def build_classification_model(
     input_specs: tf.keras.layers.InputSpec,
     model_config: classification_cfg.ImageClassificationModel,
     l2_regularizer: tf.keras.regularizers.Regularizer = None,
-    skip_logits_layer: bool = False):
+    skip_logits_layer: bool = False) -> tf.keras.Model:
   """Builds the classification model."""
   backbone = backbones.factory.build_backbone(
       input_specs=input_specs,
@@ -64,9 +64,10 @@ def build_classification_model(
   return model
 
 
-def build_maskrcnn(input_specs: tf.keras.layers.InputSpec,
+def build_maskrcnn(
+    input_specs: tf.keras.layers.InputSpec,
     model_config: maskrcnn_cfg.MaskRCNN,
-                   l2_regularizer: tf.keras.regularizers.Regularizer = None):
+    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
   """Builds Mask R-CNN model."""
   backbone = backbones.factory.build_backbone(
       input_specs=input_specs,
@@ -194,9 +195,10 @@ def build_maskrcnn(input_specs: tf.keras.layers.InputSpec,
   return model
 
 
-def build_retinanet(input_specs: tf.keras.layers.InputSpec,
+def build_retinanet(
+    input_specs: tf.keras.layers.InputSpec,
     model_config: retinanet_cfg.RetinaNet,
-                    l2_regularizer: tf.keras.regularizers.Regularizer = None):
+    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
   """Builds RetinaNet model."""
   backbone = backbones.factory.build_backbone(
       input_specs=input_specs,
@@ -253,7 +255,7 @@ def build_retinanet(input_specs: tf.keras.layers.InputSpec,
 def build_segmentation_model(
     input_specs: tf.keras.layers.InputSpec,
     model_config: segmentation_cfg.SemanticSegmentationModel,
-    l2_regularizer: tf.keras.regularizers.Regularizer = None):
+    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
   """Builds Segmentation model."""
   backbone = backbones.factory.build_backbone(
       input_specs=input_specs,

official/vision/beta/modeling/factory_3d.py

@@ -53,11 +53,12 @@ def register_model_builder(key: str):
   return registry.register(_REGISTERED_MODEL_CLS, key)
 
 
-def build_model(model_type: str,
+def build_model(
+    model_type: str,
     input_specs: tf.keras.layers.InputSpec,
     model_config: video_classification_cfg.hyperparams.Config,
     num_classes: int,
-                l2_regularizer: tf.keras.regularizers.Regularizer = None):
+    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
   """Builds backbone from a config.
 
   Args:
@@ -81,7 +82,7 @@ def build_video_classification_model(
     input_specs: tf.keras.layers.InputSpec,
     model_config: video_classification_cfg.VideoClassificationModel,
     num_classes: int,
-    l2_regularizer: tf.keras.regularizers.Regularizer = None):
+    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
   """Builds the video classification model."""
   input_specs_dict = {'image': input_specs}
   backbone = backbones.factory.build_backbone(

official/vision/beta/modeling/maskrcnn_model.py

@@ -14,6 +14,8 @@
 
 """Mask R-CNN model."""
 
+from typing import Any, Mapping, Optional, Union
+
 # Import libraries
 import tensorflow as tf
 
@@ -25,17 +27,17 @@ class MaskRCNNModel(tf.keras.Model):
   """The Mask R-CNN model."""
 
   def __init__(self,
-               backbone,
-               decoder,
-               rpn_head,
-               detection_head,
-               roi_generator,
-               roi_sampler,
-               roi_aligner,
-               detection_generator,
-               mask_head=None,
-               mask_sampler=None,
-               mask_roi_aligner=None,
+               backbone: tf.keras.Model,
+               decoder: tf.keras.Model,
+               rpn_head: tf.keras.layers.Layer,
+               detection_head: tf.keras.layers.Layer,
+               roi_generator: tf.keras.layers.Layer,
+               roi_sampler: tf.keras.layers.Layer,
+               roi_aligner: tf.keras.layers.Layer,
+               detection_generator: tf.keras.layers.Layer,
+               mask_head: Optional[tf.keras.layers.Layer] = None,
+               mask_sampler: Optional[tf.keras.layers.Layer] = None,
+               mask_roi_aligner: Optional[tf.keras.layers.Layer] = None,
                **kwargs):
     """Initializes the Mask R-CNN model.
 
@@ -85,13 +87,13 @@ class MaskRCNNModel(tf.keras.Model):
     self.mask_roi_aligner = mask_roi_aligner
 
   def call(self,
-           images,
-           image_shape,
-           anchor_boxes=None,
-           gt_boxes=None,
-           gt_classes=None,
-           gt_masks=None,
-           training=None):
+           images: tf.Tensor,
+           image_shape: tf.Tensor,
+           anchor_boxes: Optional[Mapping[str, tf.Tensor]] = None,
+           gt_boxes: tf.Tensor = None,
+           gt_classes: tf.Tensor = None,
+           gt_masks: tf.Tensor = None,
+           training: bool = None) -> Mapping[str, tf.Tensor]:
     model_outputs = {}
 
     # Feature extraction.
@@ -190,7 +192,8 @@ class MaskRCNNModel(tf.keras.Model):
     return model_outputs
 
   @property
-  def checkpoint_items(self):
+  def checkpoint_items(
+      self) -> Mapping[str, Union[tf.keras.Model, tf.keras.layers.Layer]]:
     """Returns a dictionary of items to be additionally checkpointed."""
     items = dict(
         backbone=self.backbone,
@@ -203,7 +206,7 @@ class MaskRCNNModel(tf.keras.Model):
 
     return items
 
-  def get_config(self):
+  def get_config(self) -> Mapping[str, Any]:
     return self._config_dict
 
   @classmethod

official/vision/beta/modeling/retinanet_model.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """RetinaNet."""
-from typing import List, Optional
+from typing import Any, Mapping, List, Optional, Union
 
 # Import libraries
 import tensorflow as tf
@@ -26,10 +26,10 @@ class RetinaNetModel(tf.keras.Model):
   """The RetinaNet model class."""
 
   def __init__(self,
-               backbone,
-               decoder,
-               head,
-               detection_generator,
+               backbone: tf.keras.Model,
+               decoder: tf.keras.Model,
+               head: tf.keras.layers.Layer,
+               detection_generator: tf.keras.layers.Layer,
                min_level: Optional[int] = None,
                max_level: Optional[int] = None,
                num_scales: Optional[int] = None,
@@ -74,10 +74,10 @@ class RetinaNetModel(tf.keras.Model):
     self._detection_generator = detection_generator
 
   def call(self,
-           images,
-           image_shape=None,
-           anchor_boxes=None,
-           training=None):
+           images: tf.Tensor,
+           image_shape: Optional[tf.Tensor] = None,
+           anchor_boxes: Optional[Mapping[str, tf.Tensor]] = None,
+           training: bool = None) -> Mapping[str, tf.Tensor]:
     """Forward pass of the RetinaNet model.
 
     Args:
@@ -163,7 +163,8 @@ class RetinaNetModel(tf.keras.Model):
       return outputs
 
   @property
-  def checkpoint_items(self):
+  def checkpoint_items(
+      self) -> Mapping[str, Union[tf.keras.Model, tf.keras.layers.Layer]]:
     """Returns a dictionary of items to be additionally checkpointed."""
     items = dict(backbone=self.backbone, head=self.head)
     if self.decoder is not None:
@@ -172,22 +173,22 @@ class RetinaNetModel(tf.keras.Model):
     return items
 
   @property
-  def backbone(self):
+  def backbone(self) -> tf.keras.Model:
     return self._backbone
 
   @property
-  def decoder(self):
+  def decoder(self) -> tf.keras.Model:
     return self._decoder
 
   @property
-  def head(self):
+  def head(self) -> tf.keras.layers.Layer:
     return self._head
 
   @property
-  def detection_generator(self):
+  def detection_generator(self) -> tf.keras.layers.Layer:
     return self._detection_generator
 
-  def get_config(self):
+  def get_config(self) -> Mapping[str, Any]:
     return self._config_dict
 
   @classmethod

official/vision/beta/modeling/segmentation_model.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Build segmentation models."""
+from typing import Any, Mapping, Union
 
 # Import libraries
 import tensorflow as tf
@@ -33,11 +34,8 @@ class SegmentationModel(tf.keras.Model):
   different backbones, and decoders.
   """
 
-  def __init__(self,
-               backbone,
-               decoder,
-               head,
-               **kwargs):
+  def __init__(self, backbone: tf.keras.Model, decoder: tf.keras.Model,
+               head: tf.keras.layers.Layer, **kwargs):
     """Segmentation initialization function.
 
     Args:
@@ -56,7 +54,7 @@ class SegmentationModel(tf.keras.Model):
     self.decoder = decoder
     self.head = head
 
-  def call(self, inputs, training=None):
+  def call(self, inputs: tf.Tensor, training: bool = None) -> tf.Tensor:
     backbone_features = self.backbone(inputs)
 
     if self.decoder:
@@ -67,14 +65,15 @@ class SegmentationModel(tf.keras.Model):
     return self.head(backbone_features, decoder_features)
 
   @property
-  def checkpoint_items(self):
+  def checkpoint_items(
+      self) -> Mapping[str, Union[tf.keras.Model, tf.keras.layers.Layer]]:
     """Returns a dictionary of items to be additionally checkpointed."""
     items = dict(backbone=self.backbone, head=self.head)
     if self.decoder is not None:
       items.update(decoder=self.decoder)
     return items
 
-  def get_config(self):
+  def get_config(self) -> Mapping[str, Any]:
     return self._config_dict
 
   @classmethod

official/vision/beta/modeling/video_classification_model.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """Build video classification models."""
-from typing import Mapping
+from typing import Any, Mapping, Optional, Union
 import tensorflow as tf
 
 layers = tf.keras.layers
@@ -23,15 +23,16 @@ layers = tf.keras.layers
 class VideoClassificationModel(tf.keras.Model):
   """A video classification class builder."""
 
-  def __init__(self,
+  def __init__(
+      self,
       backbone: tf.keras.Model,
       num_classes: int,
       input_specs: Mapping[str, tf.keras.layers.InputSpec] = None,
       dropout_rate: float = 0.0,
       aggregate_endpoints: bool = False,
-               kernel_initializer='random_uniform',
-               kernel_regularizer=None,
-               bias_regularizer=None,
+      kernel_initializer: str = 'random_uniform',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
       **kwargs):
     """Video Classification initialization function.
 
@@ -95,15 +96,16 @@ class VideoClassificationModel(tf.keras.Model):
         inputs=inputs, outputs=x, **kwargs)
 
   @property
-  def checkpoint_items(self):
+  def checkpoint_items(
+      self) -> Mapping[str, Union[tf.keras.Model, tf.keras.layers.Layer]]:
     """Returns a dictionary of items to be additionally checkpointed."""
     return dict(backbone=self.backbone)
 
   @property
-  def backbone(self):
+  def backbone(self) -> tf.keras.Model:
     return self._backbone
 
-  def get_config(self):
+  def get_config(self) -> Mapping[str, Any]:
     return self._config_dict
 
   @classmethod