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

Internal change

PiperOrigin-RevId: 361199986
parent 7f23d5e7
......@@ -51,14 +51,14 @@ SCALING_MAP = {
def round_repeats(repeats, multiplier, skip=False):
"""Round number of filters based on depth multiplier."""
"""Returns rounded number of filters based on depth multiplier."""
if skip or not multiplier:
return repeats
return int(math.ceil(multiplier * repeats))
def block_spec_decoder(specs, width_scale, depth_scale):
"""Decode specs for a block."""
"""Decodes and returns specs for a block."""
decoded_specs = []
for s in specs:
s = s + (
......@@ -87,7 +87,13 @@ class BlockSpec(object):
@tf.keras.utils.register_keras_serializable(package='Vision')
class EfficientNet(tf.keras.Model):
"""Class to build EfficientNet family model."""
"""Creates an EfficientNet family model.
This implements the EfficientNet model from:
Mingxing Tan, Quoc V. Le.
EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks.
(https://arxiv.org/pdf/1905.11946)
"""
def __init__(self,
model_id,
......@@ -102,25 +108,25 @@ class EfficientNet(tf.keras.Model):
norm_momentum=0.99,
norm_epsilon=0.001,
**kwargs):
"""EfficientNet initialization function.
"""Initializes an EfficientNet model.
Args:
model_id: `str` model id of EfficientNet.
input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
se_ratio: `float` squeeze and excitation ratio for inverted bottleneck
blocks.
stochastic_depth_drop_rate: `float` drop rate for drop connect layer.
kernel_initializer: kernel_initializer for convolutional layers.
kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
Default to None.
bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
model_id: A `str` of model ID of EfficientNet.
input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
se_ratio: A `float` of squeeze and excitation ratio for inverted
bottleneck blocks.
stochastic_depth_drop_rate: A `float` of drop rate for drop connect layer.
kernel_initializer: A `str` for kernel initializer of convolutional
layers.
kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
Conv2D. Default to None.
bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
Default to None.
activation: `str` name of the activation function.
use_sync_bn: if True, use synchronized batch normalization.
norm_momentum: `float` normalization omentum for the moving average.
norm_epsilon: `float` small float added to variance to avoid dividing by
zero.
**kwargs: keyword arguments to be passed.
activation: A `str` of name of the activation function.
use_sync_bn: If True, use synchronized batch normalization.
norm_momentum: A `float` of normalization momentum for the moving average.
norm_epsilon: A `float` added to variance to avoid dividing by zero.
**kwargs: Additional keyword arguments to be passed.
"""
self._model_id = model_id
self._input_specs = input_specs
......@@ -203,12 +209,12 @@ class EfficientNet(tf.keras.Model):
"""Creates one group of blocks for the EfficientNet model.
Args:
inputs: `Tensor` of size `[batch, channels, height, width]`.
specs: specifications for one inverted bottleneck block group.
name: `str`name for the block.
inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
specs: The specifications for one inverted bottleneck block group.
name: A `str` name for the block.
Returns:
The output `Tensor` of the block layer.
The output `tf.Tensor` of the block layer.
"""
if specs.block_fn == 'mbconv':
block_fn = nn_blocks.InvertedBottleneckBlock
......@@ -282,7 +288,7 @@ def build_efficientnet(
input_specs: tf.keras.layers.InputSpec,
model_config,
l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
"""Builds ResNet 3d backbone from a config."""
"""Builds EfficientNet backbone from a config."""
backbone_type = model_config.backbone.type
backbone_cfg = model_config.backbone.get()
norm_activation_config = model_config.norm_activation
......
......@@ -70,10 +70,10 @@ def register_backbone_builder(key: str):
```
Args:
key: the key to look up the builder.
key: A `str` of key to look up the builder.
Returns:
A callable for use as class decorator that registers the decorated class
A callable for using as class decorator that registers the decorated class
for creation from an instance of task_config_cls.
"""
return registry.register(_REGISTERED_BACKBONE_CLS, key)
......@@ -85,12 +85,13 @@ def build_backbone(input_specs: tf.keras.layers.InputSpec,
"""Builds backbone from a config.
Args:
input_specs: tf.keras.layers.InputSpec.
model_config: a OneOfConfig. Model config.
l2_regularizer: tf.keras.regularizers.Regularizer instance. Default to None.
input_specs: A `tf.keras.layers.InputSpec` of input.
model_config: A `OneOfConfig` of model config.
l2_regularizer: A `tf.keras.regularizers.Regularizer` object. Default to
None.
Returns:
tf.keras.Model instance of the backbone.
A `tf.keras.Model` instance of the backbone.
"""
backbone_builder = registry.lookup(_REGISTERED_BACKBONE_CLS,
model_config.backbone.type)
......
......@@ -12,9 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Contains definitions of Mobilenet Networks."""
"""Contains definitions of MobileNet Networks."""
from typing import Text, Optional, Dict, Any, Tuple
from typing import Optional, Dict, Any, Tuple
# Import libraries
import dataclasses
......@@ -41,8 +41,8 @@ class Conv2DBNBlock(tf.keras.layers.Layer):
kernel_size: int = 3,
strides: int = 1,
use_bias: bool = False,
activation: Text = 'relu6',
kernel_initializer: Text = 'VarianceScaling',
activation: str = 'relu6',
kernel_initializer: str = 'VarianceScaling',
kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
use_normalization: bool = True,
......@@ -53,25 +53,25 @@ class Conv2DBNBlock(tf.keras.layers.Layer):
"""A convolution block with batch normalization.
Args:
filters: `int` number of filters for the first two convolutions. Note that
the third and final convolution will use 4 times as many filters.
kernel_size: `int` an integer specifying the height and width of the
2D convolution window.
strides: `int` block stride. If greater than 1, this block will ultimately
downsample the input.
use_bias: if True, use biase in the convolution layer.
activation: `str` name of the activation function.
kernel_initializer: kernel_initializer for convolutional layers.
kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
filters: An `int` number of filters for the first two convolutions. Note
that the third and final convolution will use 4 times as many filters.
kernel_size: An `int` specifying the height and width of the 2D
convolution window.
strides: An `int` of block stride. If greater than 1, this block will
ultimately downsample the input.
use_bias: If True, use bias in the convolution layer.
activation: A `str` name of the activation function.
kernel_initializer: A `str` for kernel initializer of convolutional
layers.
kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
Conv2D. Default to None.
bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
Default to None.
bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
Default to None.
use_normalization: if True, use batch normalization.
use_sync_bn: if True, use synchronized batch normalization.
norm_momentum: `float` normalization momentum for the moving average.
norm_epsilon: `float` small float added to variance to avoid dividing by
zero.
**kwargs: keyword arguments to be passed.
use_normalization: If True, use batch normalization.
use_sync_bn: If True, use synchronized batch normalization.
norm_momentum: A `float` of normalization momentum for the moving average.
norm_epsilon: A `float` added to variance to avoid dividing by zero.
**kwargs: Additional keyword arguments to be passed.
"""
super(Conv2DBNBlock, self).__init__(**kwargs)
self._filters = filters
......@@ -375,13 +375,13 @@ SUPPORTED_SPECS_MAP = {
class BlockSpec(hyperparams.Config):
"""A container class that specifies the block configuration for MobileNet."""
block_fn: Text = 'convbn'
block_fn: str = 'convbn'
kernel_size: int = 3
strides: int = 1
filters: int = 32
use_bias: bool = False
use_normalization: bool = True
activation: Text = 'relu6'
activation: str = 'relu6'
# used for block type InvertedResConv
expand_ratio: Optional[float] = 6.
# used for block type InvertedResConv with SE
......@@ -395,22 +395,22 @@ def block_spec_decoder(specs: Dict[Any, Any],
# set to 1 for mobilenetv1
divisible_by: int = 8,
finegrain_classification_mode: bool = True):
"""Decode specs for a block.
"""Decodes specs for a block.
Args:
specs: `dict` specification of block specs of a mobilenet version.
filter_size_scale: `float` multiplier for the filter size
for all convolution ops. The value must be greater than zero. Typical
usage will be to set this value in (0, 1) to reduce the number of
parameters or computation cost of the model.
divisible_by: `int` ensures all inner dimensions are divisible by
specs: A `dict` specification of block specs of a mobilenet version.
filter_size_scale: A `float` multiplier for the filter size for all
convolution ops. The value must be greater than zero. Typical usage will
be to set this value in (0, 1) to reduce the number of parameters or
computation cost of the model.
divisible_by: An `int` that ensures all inner dimensions are divisible by
this number.
finegrain_classification_mode: if True, the model
will keep the last layer large even for small multipliers. Following
https://arxiv.org/abs/1801.04381
finegrain_classification_mode: If True, the model will keep the last layer
large even for small multipliers, following
https://arxiv.org/abs/1801.04381.
Returns:
List[BlockSpec]` defines structure of the base network.
A list of `BlockSpec` that defines structure of the base network.
"""
spec_name = specs['spec_name']
......@@ -449,17 +449,18 @@ def block_spec_decoder(specs: Dict[Any, Any],
@tf.keras.utils.register_keras_serializable(package='Vision')
class MobileNet(tf.keras.Model):
"""Class to build MobileNet family model."""
"""Creates a MobileNet family model."""
def __init__(self,
model_id: Text = 'MobileNetV2',
def __init__(
self,
model_id: str = 'MobileNetV2',
filter_size_scale: float = 1.0,
input_specs: layers.InputSpec = layers.InputSpec(
shape=[None, None, None, 3]),
# The followings are for hyper-parameter tuning
norm_momentum: float = 0.99,
norm_epsilon: float = 0.001,
kernel_initializer: Text = 'VarianceScaling',
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
......@@ -473,42 +474,43 @@ class MobileNet(tf.keras.Model):
# finegrain is not used in MobileNetV1
finegrain_classification_mode: bool = True,
**kwargs):
"""MobileNet initializer.
"""Initializes a MobileNet model.
Args:
model_id: `str` version of MobileNet. The supported values are
'MobileNetV1', 'MobileNetV2', 'MobileNetV3Large', 'MobileNetV3Small',
and 'MobileNetV3EdgeTPU'.
filter_size_scale: `float` multiplier for the filters (number of channels)
for all convolution ops. The value must be greater than zero. Typical
usage will be to set this value in (0, 1) to reduce the number of
parameters or computation cost of the model.
input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
norm_momentum: `float` normalization omentum for the moving average.
norm_epsilon: `float` small float added to variance to avoid dividing by
zero.
kernel_initializer: `str` kernel_initializer for convolutional layers.
kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
Default to None.
bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
model_id: A `str` of MobileNet version. The supported values are
`MobileNetV1`, `MobileNetV2`, `MobileNetV3Large`, `MobileNetV3Small`,
and `MobileNetV3EdgeTPU`.
filter_size_scale: A `float` of multiplier for the filters (number of
channels) for all convolution ops. The value must be greater than zero.
Typical usage will be to set this value in (0, 1) to reduce the number
of parameters or computation cost of the model.
input_specs: A `tf.keras.layers.InputSpec` of specs of the input tensor.
norm_momentum: A `float` of normalization momentum for the moving average.
norm_epsilon: A `float` added to variance to avoid dividing by zero.
kernel_initializer: A `str` for kernel initializer of convolutional
layers.
kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
Conv2D. Default to None.
bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
Default to None.
output_stride: `int` specifies the requested ratio of input to output
spatial resolution. If not None, then we invoke atrous convolution
if necessary to prevent the network from reducing the spatial resolution
of activation maps. Allowed values are 8 (accurate fully convolutional
mode), 16 (fast fully convolutional mode), 32 (classification mode).
min_depth: `int` minimum depth (number of channels) for all conv ops.
Enforced when filter_size_scale < 1, and not an active constraint when
filter_size_scale >= 1.
divisible_by: `int` ensures all inner dimensions are divisible by
output_stride: An `int` that specifies the requested ratio of input to
output spatial resolution. If not None, then we invoke atrous
convolution if necessary to prevent the network from reducing the
spatial resolution of activation maps. Allowed values are 8 (accurate
fully convolutional mode), 16 (fast fully convolutional mode), 32
(classification mode).
min_depth: An `int` of minimum depth (number of channels) for all
convolution ops. Enforced when filter_size_scale < 1, and not an active
constraint when filter_size_scale >= 1.
divisible_by: An `int` that ensures all inner dimensions are divisible by
this number.
stochastic_depth_drop_rate: `float` drop rate for drop connect layer.
regularize_depthwise: if Ture, apply regularization on depthwise.
use_sync_bn: if True, use synchronized batch normalization.
finegrain_classification_mode: if True, the model
will keep the last layer large even for small multipliers. Following
https://arxiv.org/abs/1801.04381
**kwargs: keyword arguments to be passed.
stochastic_depth_drop_rate: A `float` of drop rate for drop connect layer.
regularize_depthwise: If Ture, apply regularization on depthwise.
use_sync_bn: If True, use synchronized batch normalization.
finegrain_classification_mode: If True, the model will keep the last layer
large even for small multipliers, following
https://arxiv.org/abs/1801.04381.
**kwargs: Additional keyword arguments to be passed.
"""
if model_id not in SUPPORTED_SPECS_MAP:
raise ValueError('The MobileNet version {} '
......@@ -567,10 +569,10 @@ class MobileNet(tf.keras.Model):
def _mobilenet_base(self,
inputs: tf.Tensor
) -> Tuple[tf.Tensor, Dict[int, tf.Tensor]]:
"""Build the base MobileNet architecture.
"""Builds the base MobileNet architecture.
Args:
inputs: Input tensor of shape [batch_size, height, width, channels].
inputs: A `tf.Tensor` of shape `[batch_size, height, width, channels]`.
Returns:
A tuple of output Tensor and dictionary that collects endpoints.
......@@ -725,7 +727,7 @@ def build_mobilenet(
input_specs: tf.keras.layers.InputSpec,
model_config,
l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
"""Builds MobileNet 3d backbone from a config."""
"""Builds MobileNet backbone from a config."""
backbone_type = model_config.backbone.type
backbone_cfg = model_config.backbone.get()
norm_activation_config = model_config.norm_activation
......
......@@ -12,12 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Contains definitions of Residual Networks.
Residual networks (ResNets) were proposed in:
[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
Deep Residual Learning for Image Recognition. arXiv:1512.03385
"""
"""Contains definitions of Residual Networks."""
# Import libraries
import tensorflow as tf
......@@ -92,7 +87,13 @@ RESNET_SPECS = {
@tf.keras.utils.register_keras_serializable(package='Vision')
class ResNet(tf.keras.Model):
"""Class to build ResNet family model."""
"""Creates a ResNet family model.
This implements the Deep Residual Network from:
Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun.
Deep Residual Learning for Image Recognition.
(https://arxiv.org/pdf/1512.03385)
"""
def __init__(self,
model_id,
......@@ -111,32 +112,31 @@ class ResNet(tf.keras.Model):
kernel_regularizer=None,
bias_regularizer=None,
**kwargs):
"""ResNet initialization function.
"""Initializes a ResNet model.
Args:
model_id: `int` depth of ResNet backbone model.
input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
depth_multiplier: `float` a depth multiplier to uniformaly scale up all
layers in channel size in ResNet.
stem_type: `str` stem type of ResNet. Default to `v0`. If set to `v1`,
use ResNet-D type stem (https://arxiv.org/abs/1812.01187).
resnetd_shortcut: `bool` whether to use ResNet-D shortcut in downsampling
blocks.
replace_stem_max_pool: `bool` if True, replace the max pool in stem with
a stride-2 conv,
se_ratio: `float` or None. Ratio of the Squeeze-and-Excitation layer.
init_stochastic_depth_rate: `float` initial stochastic depth rate.
activation: `str` name of the activation function.
use_sync_bn: if True, use synchronized batch normalization.
norm_momentum: `float` normalization omentum for the moving average.
norm_epsilon: `float` small float added to variance to avoid dividing by
zero.
kernel_initializer: kernel_initializer for convolutional layers.
kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
Default to None.
bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
model_id: An `int` of the depth of ResNet backbone model.
input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
depth_multiplier: A `float` of the depth multiplier to uniformaly scale up
all layers in channel size in ResNet.
stem_type: A `str` of stem type of ResNet. Default to `v0`. If set to
`v1`, use ResNet-D type stem (https://arxiv.org/abs/1812.01187).
resnetd_shortcut: A `bool` of whether to use ResNet-D shortcut in
downsampling blocks.
replace_stem_max_pool: A `bool` of whether to replace the max pool in stem
with a stride-2 conv,
se_ratio: A `float` or None. Ratio of the Squeeze-and-Excitation layer.
init_stochastic_depth_rate: A `float` of initial stochastic depth rate.
activation: A `str` name of the activation function.
use_sync_bn: If True, use synchronized batch normalization.
norm_momentum: A `float` of normalization momentum for the moving average.
norm_epsilon: A small `float` added to variance to avoid dividing by zero.
kernel_initializer: A str for kernel initializer of convolutional layers.
kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
Conv2D. Default to None.
bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
Default to None.
**kwargs: keyword arguments to be passed.
**kwargs: Additional keyword arguments to be passed.
"""
self._model_id = model_id
self._input_specs = input_specs
......@@ -279,17 +279,20 @@ class ResNet(tf.keras.Model):
"""Creates one group of blocks for the ResNet model.
Args:
inputs: `Tensor` of size `[batch, channels, height, width]`.
filters: `int` number of filters for the first convolution of the layer.
strides: `int` stride to use for the first convolution of the layer. If
greater than 1, this layer will downsample the input.
block_fn: Either `nn_blocks.ResidualBlock` or `nn_blocks.BottleneckBlock`.
block_repeats: `int` number of blocks contained in the layer.
stochastic_depth_drop_rate: `float` drop rate of the current block group.
name: `str`name for the block.
inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
filters: An `int` number of filters for the first convolution of the
layer.
strides: An `int` stride to use for the first convolution of the layer.
If greater than 1, this layer will downsample the input.
block_fn: The type of block group. Either `nn_blocks.ResidualBlock` or
`nn_blocks.BottleneckBlock`.
block_repeats: An `int` number of blocks contained in the layer.
stochastic_depth_drop_rate: A `float` of drop rate of the current block
group.
name: A `str` name for the block.
Returns:
The output `Tensor` of the block layer.
The output `tf.Tensor` of the block layer.
"""
x = block_fn(
filters=filters,
......
......@@ -41,7 +41,7 @@ RESNET_SPECS = {
@tf.keras.utils.register_keras_serializable(package='Vision')
class ResNet3D(tf.keras.Model):
"""Class to build 3D ResNet family model."""
"""Creates a 3D ResNet family model."""
def __init__(self,
model_id: int,
......@@ -60,32 +60,33 @@ class ResNet3D(tf.keras.Model):
kernel_regularizer=None,
bias_regularizer=None,
**kwargs):
"""ResNet3D initialization function.
"""Initializes a 3D ResNet model.
Args:
model_id: `int` depth of ResNet backbone model.
temporal_strides: a list of integers that specifies the temporal strides
model_id: An `int` of depth of ResNet backbone model.
temporal_strides: A list of integers that specifies the temporal strides
for all 3d blocks.
temporal_kernel_sizes: a list of tuples that specifies the temporal kernel
temporal_kernel_sizes: A list of tuples that specifies the temporal kernel
sizes for all 3d blocks in different block groups.
use_self_gating: a list of booleans to specify applying self-gating module
use_self_gating: A list of booleans to specify applying self-gating module
or not in each block group. If None, self-gating is not applied.
input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
stem_conv_temporal_kernel_size: `int` temporal kernel size for the first
conv layer.
stem_conv_temporal_stride: `int` temporal stride for the first conv layer.
stem_pool_temporal_stride: `int` temporal stride for the first pool layer.
activation: `str` name of the activation function.
use_sync_bn: if True, use synchronized batch normalization.
norm_momentum: `float` normalization omentum for the moving average.
norm_epsilon: `float` small float added to variance to avoid dividing by
zero.
kernel_initializer: kernel_initializer for convolutional layers.
kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
stem_conv_temporal_kernel_size: An `int` of temporal kernel size for the
first conv layer.
stem_conv_temporal_stride: An `int` of temporal stride for the first conv
layer.
stem_pool_temporal_stride: An `int` of temporal stride for the first pool
layer.
activation: A `str` of name of the activation function.
use_sync_bn: If True, use synchronized batch normalization.
norm_momentum: A `float` of normalization momentum for the moving average.
norm_epsilon: A `float` added to variance to avoid dividing by zero.
kernel_initializer: A str for kernel initializer of convolutional layers.
kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
Conv2D. Default to None.
bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
Default to None.
bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
Default to None.
**kwargs: keyword arguments to be passed.
**kwargs: Additional keyword arguments to be passed.
"""
self._model_id = model_id
self._temporal_strides = temporal_strides
......@@ -181,21 +182,23 @@ class ResNet3D(tf.keras.Model):
"""Creates one group of blocks for the ResNet3D model.
Args:
inputs: `Tensor` of size `[batch, channels, height, width]`.
filters: `int` number of filters for the first convolution of the layer.
temporal_kernel_sizes: a tuple that specifies the temporal kernel sizes
inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
filters: An `int` of number of filters for the first convolution of the
layer.
temporal_kernel_sizes: A tuple that specifies the temporal kernel sizes
for each block in the current group.
temporal_strides: `int` temporal strides for the first convolution in this
group.
spatial_strides: `int` stride to use for the first convolution of the
temporal_strides: An `int` of temporal strides for the first convolution
in this group.
spatial_strides: An `int` stride to use for the first convolution of the
layer. If greater than 1, this layer will downsample the input.
block_fn: Either `nn_blocks.ResidualBlock` or `nn_blocks.BottleneckBlock`.
block_repeats: `int` number of blocks contained in the layer.
use_self_gating: `bool` apply self-gating module or not.
name: `str`name for the block.
block_repeats: An `int` of number of blocks contained in the layer.
use_self_gating: A `bool` that specifies whether to apply self-gating
module or not.
name: A `str` name for the block.
Returns:
The output `Tensor` of the block layer.
The output `tf.Tensor` of the block layer.
"""
if len(temporal_kernel_sizes) != block_repeats:
raise ValueError(
......
......@@ -45,12 +45,12 @@ RESNET_SPECS = {
@tf.keras.utils.register_keras_serializable(package='Vision')
class DilatedResNet(tf.keras.Model):
"""Class to build ResNet model with Deeplabv3 modifications.
"""Creates a ResNet model with Deeplabv3 modifications.
This backbone is suitable for semantic segmentation. It was proposed in:
[1] Liang-Chieh Chen, George Papandreou, Florian Schroff, Hartwig Adam
This backbone is suitable for semantic segmentation. This implements
Liang-Chieh Chen, George Papandreou, Florian Schroff, Hartwig Adam.
Rethinking Atrous Convolution for Semantic Image Segmentation.
arXiv:1706.05587
(https://arxiv.org/pdf/1706.05587)
"""
def __init__(self,
......@@ -70,30 +70,31 @@ class DilatedResNet(tf.keras.Model):
kernel_regularizer=None,
bias_regularizer=None,
**kwargs):
"""ResNet with DeepLab modification initialization function.
"""Initializes a ResNet model with DeepLab modification.
Args:
model_id: `int` depth of ResNet backbone model.
output_stride: `int` output stride, ratio of input to output resolution.
input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
stem_type: `standard` or `deeplab`, deeplab replaces 7x7 conv by 3 3x3
convs.
se_ratio: `float` or None. Ratio of the Squeeze-and-Excitation layer.
init_stochastic_depth_rate: `float` initial stochastic depth rate.
multigrid: `Tuple` of the same length as the number of blocks in the last
model_id: An `int` specifies depth of ResNet backbone model.
output_stride: An `int` of output stride, ratio of input to output
resolution.
input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
stem_type: A `str` of stem type. Can be `standard` or `deeplab`. `deeplab`
replaces 7x7 conv by 3 3x3 convs.
se_ratio: A `float` or None. Ratio of the Squeeze-and-Excitation layer.
init_stochastic_depth_rate: A `float` of initial stochastic depth rate.
multigrid: A tuple of the same length as the number of blocks in the last
resnet stage.
last_stage_repeats: `int`, how many times last stage is repeated.
activation: `str` name of the activation function.
use_sync_bn: if True, use synchronized batch normalization.
norm_momentum: `float` normalization omentum for the moving average.
norm_epsilon: `float` small float added to variance to avoid dividing by
zero.
kernel_initializer: kernel_initializer for convolutional layers.
kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
last_stage_repeats: An `int` that specifies how many times last stage is
repeated.
activation: A `str` name of the activation function.
use_sync_bn: If True, use synchronized batch normalization.
norm_momentum: A `float` of normalization momentum for the moving average.
norm_epsilon: A `float` added to variance to avoid dividing by zero.
kernel_initializer: A str for kernel initializer of convolutional layers.
kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
Conv2D. Default to None.
bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
Default to None.
bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
Default to None.
**kwargs: keyword arguments to be passed.
**kwargs: Additional keyword arguments to be passed.
"""
self._model_id = model_id
self._output_stride = output_stride
......@@ -247,20 +248,22 @@ class DilatedResNet(tf.keras.Model):
Deeplab applies strides at the last block.
Args:
inputs: `Tensor` of size `[batch, channels, height, width]`.
filters: `int` number of filters for the first convolution of the layer.
strides: `int` stride to use for the first convolution of the layer. If
greater than 1, this layer will downsample the input.
dilation_rate: `int`, diluted convolution rates.
inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
filters: An `int` off number of filters for the first convolution of the
layer.
strides: An `int` of stride to use for the first convolution of the layer.
If greater than 1, this layer will downsample the input.
dilation_rate: An `int` of diluted convolution rates.
block_fn: Either `nn_blocks.ResidualBlock` or `nn_blocks.BottleneckBlock`.
block_repeats: `int` number of blocks contained in the layer.
stochastic_depth_drop_rate: `float` drop rate of the current block group.
multigrid: List of ints or None, if specified, dilation rates for each
block_repeats: An `int` of number of blocks contained in the layer.
stochastic_depth_drop_rate: A `float` of drop rate of the current block
group.
multigrid: A list of `int` or None. If specified, dilation rates for each
block is scaled up by its corresponding factor in the multigrid.
name: `str`name for the block.
name: A `str` name for the block.
Returns:
The output `Tensor` of the block layer.
The output `tf.Tensor` of the block layer.
"""
if multigrid is not None and len(multigrid) != block_repeats:
raise ValueError('multigrid has to match number of block_repeats')
......
......@@ -13,12 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================="""
"""RevNet Implementation.
[1] Aidan N. Gomez, Mengye Ren, Raquel Urtasun, Roger B. Grosse
The Reversible Residual Network: Backpropagation Without Storing Activations
https://arxiv.org/pdf/1707.04585.pdf
"""
"""Contains definitions of RevNet."""
from typing import Any, Callable, Dict, Optional
# Import libraries
......@@ -55,7 +50,14 @@ REVNET_SPECS = {
@tf.keras.utils.register_keras_serializable(package='Vision')
class RevNet(tf.keras.Model):
"""Reversible ResNet, RevNet implementation."""
"""Creates a Reversible ResNet (RevNet) family model.
This implements:
Aidan N. Gomez, Mengye Ren, Raquel Urtasun, Roger B. Grosse.
The Reversible Residual Network: Backpropagation Without Storing
Activations.
(https://arxiv.org/pdf/1707.04585.pdf)
"""
def __init__(self,
model_id: int,
......@@ -68,19 +70,19 @@ class RevNet(tf.keras.Model):
kernel_initializer: str = 'VarianceScaling',
kernel_regularizer: tf.keras.regularizers.Regularizer = None,
**kwargs):
"""RevNet initialization function.
"""Initializes a RevNet model.
Args:
model_id: `int` depth/id of ResNet backbone model.
input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
activation: `str` name of the activation function.
use_sync_bn: `bool` if True, use synchronized batch normalization.
norm_momentum: `float` normalization omentum for the moving average.
norm_epsilon: `float` small float added to variance to avoid dividing by
zero.
kernel_initializer: `str` kernel_initializer for convolutional layers.
kernel_regularizer: `tf.keras.regularizers.Regularizer` for Conv2D.
**kwargs: additional keyword arguments to be passed.
model_id: An `int` of depth/id of ResNet backbone model.
input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
activation: A `str` name of the activation function.
use_sync_bn: If True, use synchronized batch normalization.
norm_momentum: A `float` of normalization momentum for the moving average.
norm_epsilon: A `float` added to variance to avoid dividing by zero.
kernel_initializer: A str for kernel initializer of convolutional layers.
kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
Conv2D. Default to None.
**kwargs: Additional keyword arguments to be passed.
"""
self._model_id = model_id
self._input_specs = input_specs
......@@ -148,19 +150,21 @@ class RevNet(tf.keras.Model):
"""Creates one reversible block for RevNet model.
Args:
inputs: `Tensor` of size `[batch, channels, height, width]`.
filters: `int` number of filters for the first convolution of the layer.
strides: `int` stride to use for the first convolution of the layer. If
inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
filters: An `int` number of filters for the first convolution of the
layer.
strides: An `int` stride to use for the first convolution of the layer. If
greater than 1, this block group will downsample the input.
inner_block_fn: Either `nn_blocks.ResidualInner` or
`nn_blocks.BottleneckResidualInner`.
block_repeats: `int` number of blocks contained in this block group.
batch_norm_first: `bool` whether to apply BatchNormalization and
activation layer before feeding into convolution layers.
name: `str`name for the block.
block_repeats: An `int` number of blocks contained in this block group.
batch_norm_first: A `bool` that specifies whether to apply
BatchNormalization and activation layer before feeding into convolution
layers.
name: A `str` name for the block.
Returns:
The output `Tensor` of the block layer.
The output `tf.Tensor` of the block layer.
"""
x = inputs
for i in range(block_repeats):
......@@ -210,7 +214,7 @@ def build_revnet(
input_specs: tf.keras.layers.InputSpec,
model_config,
l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
"""Builds ResNet 3d backbone from a config."""
"""Builds RevNet backbone from a config."""
backbone_type = model_config.backbone.type
backbone_cfg = model_config.backbone.get()
norm_activation_config = model_config.norm_activation
......
......@@ -13,12 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementation of SpineNet model.
X. Du, T-Y. Lin, P. Jin, G. Ghiasi, M. Tan, Y. Cui, Q. V. Le, X. Song
SpineNet: Learning Scale-Permuted Backbone for Recognition and Localization
https://arxiv.org/abs/1912.05027
"""
"""Contains definitions of SpineNet Networks."""
import math
# Import libraries
......@@ -117,7 +112,14 @@ def build_block_specs(block_specs=None):
@tf.keras.utils.register_keras_serializable(package='Vision')
class SpineNet(tf.keras.Model):
"""Class to build SpineNet models."""
"""Creates a SpineNet family model.
This implements:
Xianzhi Du, Tsung-Yi Lin, Pengchong Jin, Golnaz Ghiasi, Mingxing Tan,
Yin Cui, Quoc V. Le, Xiaodan Song.
SpineNet: Learning Scale-Permuted Backbone for Recognition and Localization.
(https://arxiv.org/abs/1912.05027)
"""
def __init__(self,
input_specs=tf.keras.layers.InputSpec(shape=[None, 640, 640, 3]),
......@@ -137,7 +139,34 @@ class SpineNet(tf.keras.Model):
norm_momentum=0.99,
norm_epsilon=0.001,
**kwargs):
"""SpineNet model."""
"""Initializes a SpineNet model.
Args:
input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
min_level: An `int` of min level for output mutiscale features.
max_level: An `int` of max level for output mutiscale features.
block_specs: The block specifications for the SpineNet model discovered by
NAS.
endpoints_num_filters: An `int` of feature dimension for the output
endpoints.
resample_alpha: A `float` of resampling factor in cross-scale connections.
block_repeats: An `int` of number of blocks contained in the layer.
filter_size_scale: A `float` of multiplier for the filters (number of
channels) for all convolution ops. The value must be greater than zero.
Typical usage will be to set this value in (0, 1) to reduce the number
of parameters or computation cost of the model.
init_stochastic_depth_rate: A `float` of initial stochastic depth rate.
kernel_initializer: A str for kernel initializer of convolutional layers.
kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
Conv2D. Default to None.
bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
Default to None.
activation: A `str` name of the activation function.
use_sync_bn: If True, use synchronized batch normalization.
norm_momentum: A `float` of normalization momentum for the moving average.
norm_epsilon: A small `float` added to variance to avoid dividing by zero.
**kwargs: Additional keyword arguments to be passed.
"""
self._input_specs = input_specs
self._min_level = min_level
self._max_level = max_level
......@@ -235,7 +264,7 @@ class SpineNet(tf.keras.Model):
return tf.identity(x, name=name)
def _build_stem(self, inputs):
"""Build SpineNet stem."""
"""Builds SpineNet stem."""
x = layers.Conv2D(
filters=64,
kernel_size=7,
......@@ -271,7 +300,7 @@ class SpineNet(tf.keras.Model):
net,
input_width,
weighted_fusion=False):
"""Build scale-permuted network."""
"""Builds scale-permuted network."""
net_sizes = [int(math.ceil(input_width / 2**2))] * len(net)
net_block_fns = [self._init_block_fn] * len(net)
num_outgoing_connections = [0] * len(net)
......@@ -363,7 +392,7 @@ class SpineNet(tf.keras.Model):
return endpoints
def _build_endpoints(self, net):
"""Match filter size for endpoints before sharing conv layers."""
"""Matches filter size for endpoints before sharing conv layers."""
endpoints = {}
for level in range(self._min_level, self._max_level + 1):
x = layers.Conv2D(
......@@ -392,7 +421,7 @@ class SpineNet(tf.keras.Model):
target_num_filters,
target_block_fn,
alpha=0.5):
"""Match resolution and feature dimension."""
"""Matches resolution and feature dimension."""
_, _, _, input_num_filters = inputs.get_shape().as_list()
if input_block_fn == 'bottleneck':
input_num_filters /= 4
......@@ -493,7 +522,7 @@ def build_spinenet(
input_specs: tf.keras.layers.InputSpec,
model_config,
l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
"""Builds ResNet 3d backbone from a config."""
"""Builds SpineNet backbone from a config."""
backbone_type = model_config.backbone.type
backbone_cfg = model_config.backbone.get()
norm_activation_config = model_config.norm_activation
......
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