diff --git a/keras_core/backend/tensorflow/layer.py b/keras_core/backend/tensorflow/layer.py index 0a96b2f4c..51f980f54 100644 --- a/keras_core/backend/tensorflow/layer.py +++ b/keras_core/backend/tensorflow/layer.py @@ -6,6 +6,25 @@ class TFLayer(tf.__internal__.tracking.AutoTrackable): """Can be overriden to perform post-build actions.""" pass + def _trackable_children(self, save_type="checkpoint", **kwargs): + if save_type == "savedmodel": + # SavedModel needs to ignore the execution functions. + train_function = getattr(self, "train_function", None) + test_function = getattr(self, "test_function", None) + predict_function = getattr(self, "predict_function", None) + self.train_function = None + self.test_function = None + self.predict_function = None + + children = super()._trackable_children(save_type, **kwargs) + + if save_type == "savedmodel": + self.train_function = train_function + self.test_function = test_function + self.predict_function = predict_function + + return children + @property def _default_save_signature(self): """For SavedModel support: returns the default serving signature.""" diff --git a/keras_core/backend/tensorflow/saved_model_test.py b/keras_core/backend/tensorflow/saved_model_test.py new file mode 100644 index 000000000..fc62503f3 --- /dev/null +++ b/keras_core/backend/tensorflow/saved_model_test.py @@ -0,0 +1,99 @@ +"""Tests for SavedModel functionality under tf implementation.""" + +import os + +import numpy as np +import pytest +import tensorflow as tf + +from keras_core import backend +from keras_core import layers +from keras_core import metrics +from keras_core import models +from keras_core import testing +from keras_core.saving import object_registration + + +@object_registration.register_keras_serializable(package="my_package") +class CustomModelX(models.Model): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.dense1 = layers.Dense(1) + self.dense2 = layers.Dense(1) + + def call(self, inputs): + out = self.dense1(inputs) + return self.dense2(out) + + def one(self): + return 1 + + +@pytest.mark.skipif( + backend.backend() != "tensorflow", + reason="The SavedModel test can only run with TF backend.", +) +class SavedModelTest(testing.TestCase): + def test_sequential(self): + model = models.Sequential([layers.Dense(1)]) + model.compile(loss="mse", optimizer="adam") + X_train = np.random.rand(100, 3) + y_train = np.random.rand(100, 1) + model.fit(X_train, y_train) + path = os.path.join(self.get_temp_dir(), "my_keras_core_model") + tf.saved_model.save(model, path) + restored_model = tf.saved_model.load(path) + self.assertAllClose( + model(X_train), + restored_model.signatures["serving_default"]( + tf.convert_to_tensor(X_train, dtype=tf.float32) + )["output_0"], + rtol=1e-4, + atol=1e-4, + ) + + def test_functional(self): + inputs = layers.Input(shape=(3,)) + x = layers.Dense(1, name="first_dense")(inputs) + outputs = layers.Dense(1, name="second_dense")(x) + model = models.Model(inputs, outputs) + model.compile( + optimizer="adam", + loss="mse", + ) + X_train = np.random.rand(100, 3) + y_train = np.random.rand(100, 1) + model.fit(X_train, y_train) + path = os.path.join(self.get_temp_dir(), "my_keras_core_model") + tf.saved_model.save(model, path) + restored_model = tf.saved_model.load(path) + self.assertAllClose( + model(X_train), + restored_model.signatures["serving_default"]( + tf.convert_to_tensor(X_train, dtype=tf.float32) + )["output_0"], + rtol=1e-4, + atol=1e-4, + ) + + def test_subclassed(self): + model = CustomModelX() + model.compile( + optimizer="adam", + loss="mse", + metrics=[metrics.Hinge(), "mse"], + ) + X_train = np.random.rand(100, 3) + y_train = np.random.rand(100, 1) + model.fit(X_train, y_train) + path = os.path.join(self.get_temp_dir(), "my_keras_core_model") + tf.saved_model.save(model, path) + restored_model = tf.saved_model.load(path) + self.assertAllClose( + model(X_train), + restored_model.signatures["serving_default"]( + tf.convert_to_tensor(X_train, dtype=tf.float32) + )["output_0"], + rtol=1e-4, + atol=1e-4, + ) diff --git a/keras_core/callbacks/learning_rate_scheduler_test.py b/keras_core/callbacks/learning_rate_scheduler_test.py index 276468f4b..568446e6a 100644 --- a/keras_core/callbacks/learning_rate_scheduler_test.py +++ b/keras_core/callbacks/learning_rate_scheduler_test.py @@ -52,8 +52,9 @@ class LearningRateSchedulerTest(testing.TestCase): lambda step: 1.0 / (1.0 + step), verbose=1 ) io_utils.disable_interactive_logging() + io_utils.set_logging_verbosity("INFO") - with self.assertLogs(level="INFO") as logs: + with self.assertLogs() as logs: self.model.fit( self.x_train, self.y_train, diff --git a/keras_core/callbacks/reduce_lr_on_plateau_test.py b/keras_core/callbacks/reduce_lr_on_plateau_test.py index 73e531dd3..015395d3a 100644 --- a/keras_core/callbacks/reduce_lr_on_plateau_test.py +++ b/keras_core/callbacks/reduce_lr_on_plateau_test.py @@ -83,8 +83,9 @@ class ReduceLROnPlateauTest(testing.TestCase): patience=1, factor=0.1, monitor="val_loss", min_delta=100, verbose=1 ) io_utils.disable_interactive_logging() + io_utils.set_logging_verbosity("INFO") - with self.assertLogs(level="INFO") as logs: + with self.assertLogs() as logs: self.model.fit( self.x_train, self.y_train, diff --git a/keras_core/layers/layer.py b/keras_core/layers/layer.py index 9b76f0c99..24060179f 100644 --- a/keras_core/layers/layer.py +++ b/keras_core/layers/layer.py @@ -295,12 +295,23 @@ class Layer(BackendLayer, Operation): ), } ) + if backend.backend() == "tensorflow": + # Remove attribute tracking for lists (TF-specific attribute) + _self_setattr_tracking = getattr( + self, "_self_setattr_tracking", True + ) + self._self_setattr_tracking = False + self._trainable_variables = trainable_variables self._non_trainable_variables = non_trainable_variables self._layers = layers self._metrics = metrics self._seed_generators = seed_generators + if backend.backend() == "tensorflow": + # Reset attribute tracking (TF-specific) + self._self_setattr_tracking = _self_setattr_tracking + @property def input_spec(self): return self._input_spec diff --git a/keras_core/utils/io_utils.py b/keras_core/utils/io_utils.py index 21127b32e..06477ce98 100644 --- a/keras_core/utils/io_utils.py +++ b/keras_core/utils/io_utils.py @@ -58,6 +58,42 @@ def is_interactive_logging_enabled(): return global_state.get_global_setting("interactive_logging", True) +@keras_core_export( + [ + "keras_core.config.set_logging_verbosity", + "keras_core.utils.set_logging_verbosity", + ] +) +def set_logging_verbosity(level): + """Sets the verbosity level for logging. + The log levels are as follows: + + - "FATAL" (least verbose) + - "ERROR" + - "WARNING" + - "INFO" + - "DEBUG" (most verbose) + + Args: + level: A string corresponding to the level of verbosity for logging. + """ + valid_levels = { + "FATAL": logging.FATAL, + "ERROR": logging.ERROR, + "WARNING": logging.WARNING, + "INFO": logging.INFO, + "DEBUG": logging.DEBUG, + } + verbosity = valid_levels.get(level) + if verbosity is None: + raise ValueError( + "Please pass a valid level for logging verbosity. " + f"The valid levels are {valid_levels.keys()}. " + f"Received: {level}" + ) + logging.set_verbosity(verbosity) + + def print_msg(message, line_break=True): """Print the message to absl logging or stdout.""" if is_interactive_logging_enabled():