Dec 5, 2025 · Dec 8, 2025 · Dec 8, 2025 · Dec 8, 2025 · Dec 8, 2025 · Dec 8, 2025
diff --git a/keras/src/callbacks/orbax_checkpoint.py b/keras/src/callbacks/orbax_checkpoint.py
 import os
 import warnings

 import numpy as np
 from keras.src.callbacks.monitor_callback import (
    MonitorCallback,  # For metric monitoring logic
 )
 from keras.src.saving.saving_lib import DiskIOStore
 from keras.src.utils.module_utils import ocp

 # Context and AsyncOptions are accessed through the lazy-loaded ocp module
                ocp.training.preservation_policies.LatestN(max_to_keep)
            )

        # Use AnyPreservationPolicy to combine them.
        # Use AnyPreservationPolicy to combine them, or use directly
        # if single policy
        preservation_policy = None
        if policies:
            preservation_policy = (
                ocp.training.preservation_policies.AnyPreservationPolicy(
                    policies
            if len(policies) == 1:
                preservation_policy = policies[0]
            else:
                preservation_policy = (
                    ocp.training.preservation_policies.AnyPreservationPolicy(
                        policies
                    )
                )
            )

        # Create the V1 Checkpointer with direct parameter passing
        # Orbax will handle directory creation on all processes as needed
            else:
                self.checkpointer.save_pytree(step, composite_state)

        # Save assets separately since PyTree can't handle binary data
        if not self.save_weights_only:
            self._save_assets(step)

    def _save_assets(self, step):
        """Save model assets to a separate directory."""
        from keras.src.saving.saving_lib import _save_state

        assets_dir = os.path.join(self.directory, "assets", str(step))
        try:
            assets_store = DiskIOStore(assets_dir, mode="w")
        except FileExistsError:
            # Directory already exists, skip asset saving
            return
        try:
            # Use the same recursive saving logic as _save_state
            visited = set()
            _save_state(
                self.model,
                None,  # No weights store
                assets_store,  # Assets store
                "",  # Root path
                visited,
            )
        finally:
            assets_store.close()

    def on_train_batch_end(self, batch, logs=None):
        if self._should_save_on_batch(batch):
            # Handle save_best_only logic for batch-level saving
Original file line number	Diff line number	Diff line change
		@@ -1,3 +1,4 @@
		import os
		import warnings

		import numpy as np
Expand All		@@ -8,6 +9,7 @@
		from keras.src.callbacks.monitor_callback import (
		MonitorCallback, # For metric monitoring logic
		)
		from keras.src.saving.saving_lib import DiskIOStore
		from keras.src.utils.module_utils import ocp

		# Context and AsyncOptions are accessed through the lazy-loaded ocp module
Expand DownExpand Up		@@ -158,14 +160,18 @@ def __init__(
		ocp.training.preservation_policies.LatestN(max_to_keep)
		)

		# Use AnyPreservationPolicy to combine them.
		# Use AnyPreservationPolicy to combine them, or use directly
		# if single policy
		preservation_policy = None
		if policies:
		preservation_policy = (
		ocp.training.preservation_policies.AnyPreservationPolicy(
		policies
		if len(policies) == 1:
		preservation_policy = policies[0]
		else:
		preservation_policy = (
		ocp.training.preservation_policies.AnyPreservationPolicy(
		policies
		)
		)
		)

		# Create the V1 Checkpointer with direct parameter passing
		# Orbax will handle directory creation on all processes as needed
Expand DownExpand Up		@@ -270,6 +276,33 @@ def _save_checkpoint(self, step, logs=None):
		else:
		self.checkpointer.save_pytree(step, composite_state)

		# Save assets separately since PyTree can't handle binary data
		if not self.save_weights_only:
		self._save_assets(step)

		def _save_assets(self, step):
		"""Save model assets to a separate directory."""
		from keras.src.saving.saving_lib import _save_state

		assets_dir = os.path.join(self.directory, "assets", str(step))
amitsrivastava78 marked this conversation as resolved. OutdatedShow resolvedHide resolved
		try:
		assets_store = DiskIOStore(assets_dir, mode="w")
		except FileExistsError:
		# Directory already exists, skip asset saving
		return
		try:
		# Use the same recursive saving logic as _save_state
		visited = set()
		_save_state(
		self.model,
		None, # No weights store
		assets_store, # Assets store
		"", # Root path
		visited,
		)
		finally:
		assets_store.close()

		def on_train_batch_end(self, batch, logs=None):
		if self._should_save_on_batch(batch):
		# Handle save_best_only logic for batch-level saving
Expand Down