Add max_shard_size to Backbone and Task. Simplify the test.

Author: james77777778

keras_hub/src/models/backbone.py

@@ -177,14 +177,17 @@ class like `keras_hub.models.Backbone.from_preset()`, or from
             )
         return loader.load_backbone(backbone_cls, load_weights, **kwargs)
 
-    def save_to_preset(self, preset_dir):
+    def save_to_preset(self, preset_dir, max_shard_size=10):
         """Save backbone to a preset directory.
 
         Args:
             preset_dir: The path to the local model preset directory.
+            max_shard_size: `int` or `float`. Maximum size in GB for each
+                sharded file. If `None`, no sharding will be done. Defaults to
+                `10`.
         """
         saver = get_preset_saver(preset_dir)
-        saver.save_backbone(self)
+        saver.save_backbone(self, max_shard_size=max_shard_size)
 
     def get_lora_target_names(self):
         """Returns list of layer names which are to be LoRA-fied.

keras_hub/src/models/task.py

@@ -236,14 +236,17 @@ def save_task_weights(self, filepath):
             objects_to_skip=backbone_layer_ids,
         )
 
-    def save_to_preset(self, preset_dir):
+    def save_to_preset(self, preset_dir, max_shard_size=10):
         """Save task to a preset directory.
 
         Args:
             preset_dir: The path to the local model preset directory.
+            max_shard_size: `int` or `float`. Maximum size in GB for each
+                sharded file. If `None`, no sharding will be done. Defaults to
+                `10`.
         """
         saver = get_preset_saver(preset_dir)
-        saver.save_task(self)
+        saver.save_task(self, max_shard_size=max_shard_size)
 
     @property
     def layers(self):

keras_hub/src/utils/preset_utils.py

@@ -1,6 +1,5 @@
 import collections
 import datetime
-import functools
 import inspect
 import json
 import math
@@ -804,18 +803,20 @@ def save_audio_converter(self, converter):
     def save_image_converter(self, converter):
         self._save_serialized_object(converter, IMAGE_CONVERTER_CONFIG_FILE)
 
-    def save_task(self, task):
+    def save_task(self, task, max_shard_size=10):
         # Save task specific config and weights.
         self._save_serialized_object(task, TASK_CONFIG_FILE)
         if task.has_task_weights():
             task_weight_path = os.path.join(self.preset_dir, TASK_WEIGHTS_FILE)
             task.save_task_weights(task_weight_path)
         # Save backbone.
         if hasattr(task.backbone, "save_to_preset"):
-            task.backbone.save_to_preset(self.preset_dir)
+            task.backbone.save_to_preset(
+                self.preset_dir, max_shard_size=max_shard_size
+            )
         else:
             # Allow saving a `keras.Model` that is not a backbone subclass.
-            self.save_backbone(task.backbone)
+            self.save_backbone(task.backbone, max_shard_size=max_shard_size)
         # Save preprocessor.
         if task.preprocessor and hasattr(task.preprocessor, "save_to_preset"):
             task.preprocessor.save_to_preset(self.preset_dir)
@@ -874,20 +875,17 @@ def _save_metadata(self, layer):
             metadata_file.write(json.dumps(metadata, indent=4))
 
     def _get_variables_size_in_bytes(self, variables):
-        @functools.lru_cache(512)
         def _compute_memory_size(shape, dtype):
+            def _get_dtype_size(dtype):
+                dtype = keras.backend.standardize_dtype(dtype)
+                # If dtype is bool, return 1 immediately.
+                if dtype == "bool":
+                    return 1
+                # Else, we extract the bit size from the string.
+                return int(re.sub(r"bfloat|float|uint|int", "", dtype))
+
             weight_counts = math.prod(shape)
-            dtype = keras.backend.standardize_dtype(dtype)
-            dtype_size = int(
-                (
-                    dtype.replace("bfloat", "")
-                    .replace("float", "")
-                    .replace("uint", "")
-                    .replace("int", "")
-                    .replace("bool", "1")
-                )
-            )
-            return weight_counts * dtype_size
+            return weight_counts * _get_dtype_size(dtype)
 
         unique_variables = {}
         for v in variables:

keras_hub/src/utils/preset_utils_test.py

@@ -14,7 +14,6 @@
 from keras_hub.src.tests.test_case import TestCase
 from keras_hub.src.utils.keras_utils import sharded_weights_available
 from keras_hub.src.utils.preset_utils import CONFIG_FILE
-from keras_hub.src.utils.preset_utils import get_preset_saver
 from keras_hub.src.utils.preset_utils import upload_preset
 
 
@@ -26,13 +25,13 @@ def test_sharded_weights(self):
 
         # Gemma2 config.
         init_kwargs = {
-            "vocabulary_size": 4096,  # 256128
-            "num_layers": 24,  # 46
-            "num_query_heads": 16,  # 32
-            "num_key_value_heads": 8,  # 16
-            "hidden_dim": 64,  # 4608
-            "intermediate_dim": 128,  # 73728
-            "head_dim": 8,  # 128
+            "vocabulary_size": 1024,  # 256128
+            "num_layers": 12,  # 46
+            "num_query_heads": 8,  # 32
+            "num_key_value_heads": 4,  # 16
+            "hidden_dim": 32,  # 4608
+            "intermediate_dim": 64,  # 73728
+            "head_dim": 4,  # 128
             "sliding_window_size": 5,  # 4096
             "attention_logit_soft_cap": 50,
             "final_logit_soft_cap": 30,
@@ -42,12 +41,12 @@ def test_sharded_weights(self):
             "use_post_attention_norm": True,
             "use_sliding_window_attention": True,
         }
-        backbone = GemmaBackbone(**init_kwargs)  # ~4.4MB
+        backbone = GemmaBackbone(**init_kwargs)  # ~422KB
+        backbone.summary()
 
         # Save the sharded weights.
         preset_dir = self.get_temp_dir()
-        preset_saver = get_preset_saver(preset_dir)
-        preset_saver.save_backbone(backbone, max_shard_size=0.002)
+        backbone.save_to_preset(preset_dir, max_shard_size=0.0002)
         self.assertTrue(
             os.path.exists(os.path.join(preset_dir, "model.weights.json"))
         )