Skip to content

GitLab

Commit f8f4845c authored by Fan Yang's avatar Fan Yang Committed by A. Unique TensorFlower
Browse files

Add more pytype checking. 

PiperOrigin-RevId: 368129317
parent 4334a892
Hide whitespace changes
Inline Side-by-side
Showing with 424 additions and 352 deletions
official/vision/beta/modeling/backbones/efficientnet.py
View file @ f8f4845c
......@@ -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
View file @ f8f4845c
......@@ -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
View file @ f8f4845c
......@@ -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,20 +416,21 @@ 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],
filter_size_scale: float,
# set to 1 for mobilenetv1
divisible_by: int = 8,
finegrain_classification_mode: bool = True):
def block_spec_decoder(
specs: Dict[Any, Any],
filter_size_scale: float,
# Set to 1 for mobilenetv1.
divisible_by: int = 8,
finegrain_classification_mode: bool = True):
"""Decodes specs for a block.
Args:
......@@ -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
View file @ f8f4845c
......@@ -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,23 +103,25 @@ 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,
**kwargs):
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.
Args:
......@@ -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
View file @ f8f4845c
......@@ -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,26 +76,28 @@ RESNET_SPECS = {
class ResNet3D(tf.keras.Model):
"""Creates a 3D ResNet family model."""
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,
**kwargs):
def __init__(
self,
model_id: int,
temporal_strides: List[int],
temporal_kernel_sizes: List[Tuple[int]],
use_self_gating: List[int] = 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.
Args:
......@@ -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
View file @ f8f4845c
......@@ -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,23 +55,25 @@ 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,
**kwargs):
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.
Args:
......@@ -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
View file @ f8f4845c
......@@ -59,17 +59,18 @@ class RevNet(tf.keras.Model):
(https://arxiv.org/pdf/1707.04585.pdf)
"""
def __init__(self,
model_id: int,
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,
**kwargs):
def __init__(
self,
model_id: int,
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: Optional[tf.keras.regularizers.Regularizer] = None,
**kwargs):
"""Initializes a RevNet model.
Args:
......
official/vision/beta/modeling/backbones/spinenet.py
View file @ f8f4845c
......@@ -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,32 +126,34 @@ 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,
**kwargs):
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.
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.
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
View file @ f8f4845c
......@@ -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,25 +131,27 @@ 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,
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,
**kwargs):
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: 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.
Args:
......@@ -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
View file @ f8f4845c
......@@ -14,6 +14,7 @@
"""Build classification models."""
from typing import Any, Mapping, Optional
# Import libraries
import tensorflow as tf
......@@ -24,20 +25,22 @@ 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,
use_sync_bn: bool = False,
norm_momentum: float = 0.99,
norm_epsilon: float = 0.001,
skip_logits_layer: bool = False,
**kwargs):
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,
skip_logits_layer: bool = False,
**kwargs):
"""Classification initialization function.
Args:
......@@ -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
View file @ f8f4845c
......@@ -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,20 +25,21 @@ 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',
**kwargs):
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.
Args:
......@@ -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
View file @ f8f4845c
......@@ -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
View file @ f8f4845c
......@@ -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,20 +33,21 @@ 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,
**kwargs):
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).
Args:
......@@ -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
View file @ f8f4845c
......@@ -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,22 +66,23 @@ 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,
**kwargs):
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.
Args:
......@@ -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
View file @ f8f4845c
......@@ -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,
model_config: maskrcnn_cfg.MaskRCNN,
l2_regularizer: tf.keras.regularizers.Regularizer = None):
def build_maskrcnn(
input_specs: tf.keras.layers.InputSpec,
model_config: maskrcnn_cfg.MaskRCNN,
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,
model_config: retinanet_cfg.RetinaNet,
l2_regularizer: tf.keras.regularizers.Regularizer = None):
def build_retinanet(
input_specs: tf.keras.layers.InputSpec,
model_config: retinanet_cfg.RetinaNet,
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
View file @ f8f4845c
......@@ -53,11 +53,12 @@ def register_model_builder(key: str):
return registry.register(_REGISTERED_MODEL_CLS, key)
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):
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) -> 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
View file @ f8f4845c
......@@ -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
View file @ f8f4845c
......@@ -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
View file @ f8f4845c
......@@ -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
View file @ f8f4845c
......@@ -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,16 +23,17 @@ layers = tf.keras.layers
class VideoClassificationModel(tf.keras.Model):
"""A video classification class builder."""
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,
**kwargs):
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: str = 'random_uniform',
kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
**kwargs):
"""Video Classification initialization function.
Args:
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment