Temporary update the registration package name.

PiperOrigin-RevId: 427335324

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."""
 

official/vision/beta/modeling/retinanet_model.py

@@ -21,7 +21,7 @@ import tensorflow as tf
 from official.vision.beta.ops import anchor
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class RetinaNetModel(tf.keras.Model):
   """The RetinaNet model class."""
 

official/vision/beta/modeling/segmentation_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 SegmentationModel(tf.keras.Model):
   """A Segmentation class model.
 

official/vision/beta/modeling/video_classification_model.py

@@ -20,7 +20,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 VideoClassificationModel(tf.keras.Model):
   """A video classification class builder."""