Temporary update the registration package name.

PiperOrigin-RevId: 427335324

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

@@ -92,7 +92,7 @@ def block_spec_decoder(specs: List[Tuple[Any, ...]], width_scale: float,
   return decoded_specs
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class EfficientNet(tf.keras.Model):
   """Creates an EfficientNet family model.
 

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

@@ -344,7 +344,7 @@ def block_spec_decoder(
   return decoded_specs
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class MobileDet(tf.keras.Model):
   """Creates a MobileDet family model."""
 

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

@@ -31,7 +31,7 @@ layers = tf.keras.layers
 #  pylint: disable=pointless-string-statement
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class Conv2DBNBlock(tf.keras.layers.Layer):
   """A convolution block with batch normalization."""
 
@@ -616,7 +616,7 @@ def block_spec_decoder(
   return decoded_specs
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class MobileNet(tf.keras.Model):
   """Creates a MobileNet family model."""
 

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

@@ -95,7 +95,7 @@ RESNET_SPECS = {
 }
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class ResNet(tf.keras.Model):
   """Creates ResNet and ResNet-RS family models.
 

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

@@ -72,7 +72,7 @@ RESNET_SPECS = {
 }
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class ResNet3D(tf.keras.Model):
   """Creates a 3D ResNet family model."""
 

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

@@ -46,7 +46,7 @@ RESNET_SPECS = {
 }
 
 
-@tf.keras.utils.register_keras_serializable(package='Vision')
+@tf.keras.utils.register_keras_serializable(package='Beta')
 class DilatedResNet(tf.keras.Model):
   """Creates a ResNet model with Deeplabv3 modifications.
 

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