Merge branch 'tensorflow:master' into panoptic-deeplab-modeling

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

@@ -49,7 +49,7 @@ REVNET_SPECS = {
 }
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class RevNet(tf.keras.Model):
   """Creates a Reversible ResNet (RevNet) family model.
 

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

@@ -124,7 +124,7 @@ def build_block_specs(
   return [BlockSpec(*b) for b in block_specs]
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class SpineNet(tf.keras.Model):
   """Creates a SpineNet family model.
 

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

@@ -117,7 +117,7 @@ def build_block_specs(
   return [BlockSpec(*b) for b in block_specs]
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class SpineNetMobile(tf.keras.Model):
   """Creates a Mobile SpineNet family model.
 

official/vision/beta/modeling/classification_model.py

@@ -21,7 +21,7 @@ import tensorflow as tf
 layers = tf.keras.layers
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class ClassificationModel(tf.keras.Model):
   """A classification class builder."""
 

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

@@ -27,7 +27,7 @@ from official.vision.beta.modeling.layers import nn_layers
 TensorMapUnion = Union[tf.Tensor, Mapping[str, tf.Tensor]]
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class ASPP(tf.keras.layers.Layer):
   """Creates an Atrous Spatial Pyramid Pooling (ASPP) layer."""
 

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

@@ -25,7 +25,7 @@ from official.vision.beta.modeling.decoders import factory
 from official.vision.beta.ops import spatial_transform_ops
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class FPN(tf.keras.Model):
   """Creates a Feature Pyramid Network (FPN).
 

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

@@ -61,7 +61,7 @@ def build_block_specs(
   return [BlockSpec(*b) for b in block_specs]
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class NASFPN(tf.keras.Model):
   """Creates a NAS-FPN model.
 

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

@@ -24,7 +24,7 @@ import tensorflow as tf
 from official.modeling import tf_utils
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class RetinaNetHead(tf.keras.layers.Layer):
   """Creates a RetinaNet head."""
 
@@ -336,7 +336,7 @@ class RetinaNetHead(tf.keras.layers.Layer):
     return cls(**config)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class RPNHead(tf.keras.layers.Layer):
   """Creates a Region Proposal Network (RPN) head."""
 

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

@@ -21,7 +21,7 @@ import tensorflow as tf
 from official.modeling import tf_utils
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class DetectionHead(tf.keras.layers.Layer):
   """Creates a detection head."""
 
@@ -216,7 +216,7 @@ class DetectionHead(tf.keras.layers.Layer):
     return cls(**config)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class MaskHead(tf.keras.layers.Layer):
   """Creates a mask head."""
 

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

@@ -191,7 +191,7 @@ class MaskScoring(tf.keras.Model):
     return cls(**config)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class SegmentationHead(tf.keras.layers.Layer):
   """Creates a segmentation head."""
 

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

@@ -20,7 +20,7 @@ import tensorflow as tf
 from official.vision.beta.ops import sampling_ops
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class BoxSampler(tf.keras.layers.Layer):
   """Creates a BoxSampler to sample positive and negative boxes."""
 

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

@@ -423,7 +423,7 @@ def _generate_detections_batched(boxes: tf.Tensor, scores: tf.Tensor,
   return nmsed_boxes, nmsed_scores, nmsed_classes, valid_detections
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class DetectionGenerator(tf.keras.layers.Layer):
   """Generates the final detected boxes with scores and classes."""
 
@@ -606,7 +606,7 @@ class DetectionGenerator(tf.keras.layers.Layer):
     return cls(**config)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class MultilevelDetectionGenerator(tf.keras.layers.Layer):
   """Generates detected boxes with scores and classes for one-stage detector."""
 

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

@@ -100,7 +100,7 @@ def _sample_and_crop_foreground_masks(candidate_rois: tf.Tensor,
   return foreground_rois, foreground_classes, cropped_foreground_masks
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class MaskSampler(tf.keras.layers.Layer):
   """Samples and creates mask training targets."""
 

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

@@ -52,7 +52,7 @@ def _maybe_downsample(x: tf.Tensor, out_filter: int, strides: int,
   return x + 0.
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class ResidualBlock(tf.keras.layers.Layer):
   """A residual block."""
 
@@ -250,7 +250,7 @@ class ResidualBlock(tf.keras.layers.Layer):
     return self._activation_fn(x + shortcut)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class BottleneckBlock(tf.keras.layers.Layer):
   """A standard bottleneck block."""
 
@@ -479,7 +479,7 @@ class BottleneckBlock(tf.keras.layers.Layer):
     return self._activation3(x)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class InvertedBottleneckBlock(tf.keras.layers.Layer):
   """An inverted bottleneck block."""
 
@@ -754,7 +754,7 @@ class InvertedBottleneckBlock(tf.keras.layers.Layer):
     return x
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class ResidualInner(tf.keras.layers.Layer):
   """Creates a single inner block of a residual.
 
@@ -878,7 +878,7 @@ class ResidualInner(tf.keras.layers.Layer):
     return x
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class BottleneckResidualInner(tf.keras.layers.Layer):
   """Creates a single inner block of a bottleneck.
 
@@ -1018,7 +1018,7 @@ class BottleneckResidualInner(tf.keras.layers.Layer):
     return x
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class ReversibleLayer(tf.keras.layers.Layer):
   """Creates a reversible layer.
 
@@ -1194,7 +1194,7 @@ class ReversibleLayer(tf.keras.layers.Layer):
     return activations
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class DepthwiseSeparableConvBlock(tf.keras.layers.Layer):
   """Creates an depthwise separable convolution block with batch normalization."""
 
@@ -1319,7 +1319,7 @@ class DepthwiseSeparableConvBlock(tf.keras.layers.Layer):
     return self._activation_fn(x)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class TuckerConvBlock(tf.keras.layers.Layer):
   """An Tucker block (generalized bottleneck)."""
 

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

@@ -20,7 +20,7 @@ from official.modeling import tf_utils
 from official.vision.beta.modeling.layers import nn_layers
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class SelfGating(tf.keras.layers.Layer):
   """Feature gating as used in S3D-G.
 
@@ -67,7 +67,7 @@ class SelfGating(tf.keras.layers.Layer):
     return tf.math.multiply(x, inputs)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class BottleneckBlock3D(tf.keras.layers.Layer):
   """Creates a 3D bottleneck block."""
 

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

@@ -85,7 +85,7 @@ def get_padding_for_kernel_size(kernel_size):
         kernel_size))
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class SqueezeExcitation(tf.keras.layers.Layer):
   """Creates a squeeze and excitation layer."""
 
@@ -223,7 +223,7 @@ def get_stochastic_depth_rate(init_rate, i, n):
   return rate
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class StochasticDepth(tf.keras.layers.Layer):
   """Creates a stochastic depth layer."""
 
@@ -261,7 +261,7 @@ class StochasticDepth(tf.keras.layers.Layer):
     return output
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 def pyramid_feature_fusion(inputs, target_level):
   """Fuses all feature maps in the feature pyramid at the target level.
 
@@ -403,7 +403,7 @@ class PanopticFPNFusion(tf.keras.Model):
     return self._output_specs
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class Scale(tf.keras.layers.Layer):
   """Scales the input by a trainable scalar weight.
 
@@ -456,7 +456,7 @@ class Scale(tf.keras.layers.Layer):
     return scale * inputs
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class TemporalSoftmaxPool(tf.keras.layers.Layer):
   """Creates a network layer corresponding to temporal softmax pooling.
 
@@ -479,7 +479,7 @@ class TemporalSoftmaxPool(tf.keras.layers.Layer):
     return outputs
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class PositionalEncoding(tf.keras.layers.Layer):
   """Creates a network layer that adds a sinusoidal positional encoding.
 
@@ -647,7 +647,7 @@ class PositionalEncoding(tf.keras.layers.Layer):
     return (outputs, states) if output_states else outputs
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class GlobalAveragePool3D(tf.keras.layers.Layer):
   """Creates a global average pooling layer with causal mode.
 
@@ -781,7 +781,7 @@ class GlobalAveragePool3D(tf.keras.layers.Layer):
     return (x, states) if output_states else x
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class SpatialAveragePool3D(tf.keras.layers.Layer):
   """Creates a global average pooling layer pooling across spatial dimentions."""
 
@@ -901,7 +901,7 @@ class CausalConvMixin:
     return spatial_output_shape
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class Conv2D(tf.keras.layers.Conv2D, CausalConvMixin):
   """Conv2D layer supporting CausalConv.
 
@@ -949,7 +949,7 @@ class Conv2D(tf.keras.layers.Conv2D, CausalConvMixin):
     return self._buffered_spatial_output_shape(shape)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class DepthwiseConv2D(tf.keras.layers.DepthwiseConv2D, CausalConvMixin):
   """DepthwiseConv2D layer supporting CausalConv.
 
@@ -1011,7 +1011,7 @@ class DepthwiseConv2D(tf.keras.layers.DepthwiseConv2D, CausalConvMixin):
     return self._buffered_spatial_output_shape(shape)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class Conv3D(tf.keras.layers.Conv3D, CausalConvMixin):
   """Conv3D layer supporting CausalConv.
 
@@ -1067,7 +1067,7 @@ class Conv3D(tf.keras.layers.Conv3D, CausalConvMixin):
     return self._buffered_spatial_output_shape(shape)
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class SpatialPyramidPooling(tf.keras.layers.Layer):
   """Implements the Atrous Spatial Pyramid Pooling.
 

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

@@ -20,7 +20,7 @@ import tensorflow as tf
 from official.vision.beta.ops import spatial_transform_ops
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class MultilevelROIAligner(tf.keras.layers.Layer):
   """Performs ROIAlign for the second stage processing."""
 

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

@@ -176,7 +176,7 @@ def _multilevel_propose_rois(raw_boxes: Mapping[str, tf.Tensor],
     return selected_rois, selected_roi_scores
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class MultilevelROIGenerator(tf.keras.layers.Layer):
   """Proposes RoIs for the second stage processing."""
 

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

@@ -23,7 +23,7 @@ from official.vision.beta.ops import iou_similarity
 from official.vision.beta.ops import target_gather
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class ROISampler(tf.keras.layers.Layer):
   """Samples ROIs and assigns targets to the sampled ROIs."""
 

official/vision/beta/modeling/maskrcnn_model.py

@@ -22,7 +22,7 @@ from official.vision.beta.ops import anchor
 from official.vision.beta.ops import box_ops
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class MaskRCNNModel(tf.keras.Model):
   """The Mask R-CNN(-RS) and Cascade RCNN-RS models."""