refactoring

DavidLandup0 · Aug 18, 2023 · ab10136 · ab10136
1 parent f882b3e
commit ab10136
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 5 deletions.
diff --git a/keras_cv/layers/hierarchical_transformer_encoder.py b/keras_cv/layers/hierarchical_transformer_encoder.py
@@ -8,7 +8,7 @@
 from keras_cv.layers.regularization.drop_path import DropPath
 
 
-@keras.saving.register_keras_serializable(package="keras_cv")
+@keras_cv_export("keras_cv.layers.HierarchicalTransformerEncoder")
 class HierarchicalTransformerEncoder(keras.layers.Layer):
     """
     Hierarchical transformer encoder block implementation as a Keras Layer.

diff --git a/keras_cv/layers/overlapping_patching_embedding.py b/keras_cv/layers/overlapping_patching_embedding.py
@@ -2,7 +2,7 @@
 from keras_cv.backend import ops
 
 
-@keras.saving.register_keras_serializable(package="keras_cv")
+@keras_cv_export("keras_cv.layers.OverlappingPatchingAndEmbedding")
 class OverlappingPatchingAndEmbedding(keras.layers.Layer):
     def __init__(self, project_dim=32, patch_size=7, stride=4, **kwargs):
         """

diff --git a/...v/layers/efficient_multihead_attention.py → ...v/layers/segformer_multihead_attention.py b/...v/layers/efficient_multihead_attention.py → ...v/layers/segformer_multihead_attention.py
@@ -4,7 +4,7 @@
 from keras_cv.backend import ops
 
 
-@keras.saving.register_keras_serializable(package="keras_cv")
+@keras_cv_export("keras_cv.layers.SegFormerMultiheadAttention")
 class SegFormerMultiheadAttention(keras.layers.Layer):
     def __init__(self, project_dim, num_heads, sr_ratio):
         """

diff --git a/keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py b/keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py
@@ -39,7 +39,7 @@
 from keras_cv.utils.python_utils import classproperty
 
 
-@keras.saving.register_keras_serializable(package="keras_cv.models")
+@keras_cv_export("keras_cv.layers.MiTBackbone")
 class MiTBackbone(Backbone):
     def __init__(
         self,

diff --git a/keras_cv/models/segmentation/segformer/segformer.py b/keras_cv/models/segmentation/segformer/segformer.py
@@ -12,7 +12,7 @@
 from keras_cv.utils.train import get_feature_extractor
 
 
-@keras.utils.register_keras_serializable(package="keras_cv")
+@keras_cv_export("keras_cv.layers.SegFormer")
 class SegFormer(Task):
     """A Keras model implementing the SegFormer architecture for semantic
     segmentation.