diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 611247e2..bf499c4f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -9,10 +9,17 @@ jobs: strategy: fail-fast: false matrix: + python-verison: ["3.7"] ml-deps: - "torch==1.11.0+cpu torchvision==0.12.0+cpu torchdata==0.3.0 tensorflow-cpu==2.8.1" - "torch==1.12.1+cpu torchvision==0.13.1+cpu torchdata==0.4.1 tensorflow-cpu==2.9.1" - "torch==1.13.0+cpu torchvision==0.14.0+cpu torchdata==0.5.0 tensorflow-cpu==2.10.0" + - "torch==1.13.0+cpu torchvision==0.14.0+cpu torchdata==0.5.0 tensorflow-cpu==2.11.0" + include: + - ml-deps: "torch==1.13.0+cpu torchvision==0.14.0+cpu torchdata==0.5.0 tensorflow-cpu==2.12.0" + python-version: "3.9" + - ml-deps: "torch==1.13.0+cpu torchvision==0.14.0+cpu torchdata==0.5.0 tensorflow-cpu==2.13.0" + python-version: "3.9" env: run_coverage: ${{ github.ref == 'refs/heads/master' }} @@ -23,7 +30,7 @@ jobs: - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v4 with: - python-version: "3.7" + python-version: ${{ matrix.python-version }} - name: Cache dependencies uses: actions/cache@v3 diff --git a/tests/models/test_tensorflow_keras_models.py b/tests/models/test_tensorflow_keras_models.py index 73a07c78..36d601bc 100644 --- a/tests/models/test_tensorflow_keras_models.py +++ b/tests/models/test_tensorflow_keras_models.py @@ -21,7 +21,16 @@ get_small_sequential_mlp, ) except ImportError: - from keras.testing_utils import get_small_functional_mlp, get_small_sequential_mlp + try: + from keras.testing_utils import ( + get_small_functional_mlp, + get_small_sequential_mlp, + ) + except ImportError: + from keras.src.testing_infra.test_utils import ( + get_small_functional_mlp, + get_small_sequential_mlp, + ) # Suppress all Tensorflow messages os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" @@ -166,8 +175,20 @@ def test_save_model_to_tiledb_array_weights( data = np.random.rand(100, 3) if optimizer: - model_opt_weights = batch_get_value(model.optimizer.weights) - loaded_opt_weights = batch_get_value(loaded_model.optimizer.weights) + if hasattr(model.optimizer, "weights"): + model_opt_weights = tf.keras.backend.batch_get_value( + model.optimizer.weights + ) + else: + model_opt_weights = [var.numpy() for var in model.optimizer.variables()] + if hasattr(loaded_model.optimizer, "weights"): + loaded_opt_weights = tf.keras.backend.batch_get_value( + loaded_model.optimizer.weights + ) + else: + loaded_opt_weights = [ + var.numpy() for var in loaded_model.optimizer.variables() + ] # Assert optimizer weights are equal for weight_model, weight_loaded_model in zip( @@ -209,8 +230,20 @@ def test_save_load_with_dense_features(self, tmpdir, loss, optimizer, metrics): tiledb_model_obj.save(include_optimizer=True) loaded_model = tiledb_model_obj.load(compile_model=True) - model_opt_weights = batch_get_value(model.optimizer.weights) - loaded_opt_weights = batch_get_value(loaded_model.optimizer.weights) + if hasattr(model.optimizer, "weights"): + model_opt_weights = tf.keras.backend.batch_get_value( + model.optimizer.weights + ) + else: + model_opt_weights = [var.numpy() for var in model.optimizer.variables()] + if hasattr(loaded_model.optimizer, "weights"): + loaded_opt_weights = tf.keras.backend.batch_get_value( + loaded_model.optimizer.weights + ) + else: + loaded_opt_weights = [ + var.numpy() for var in loaded_model.optimizer.variables() + ] # Assert optimizer weights are equal for weight_model, weight_loaded_model in zip( @@ -260,8 +293,20 @@ def test_save_load_with_sequence_features(self, tmpdir, loss, optimizer, metrics tiledb_model_obj.save(include_optimizer=True) loaded_model = tiledb_model_obj.load(compile_model=True) - model_opt_weights = batch_get_value(model.optimizer.weights) - loaded_opt_weights = batch_get_value(loaded_model.optimizer.weights) + if hasattr(model.optimizer, "weights"): + model_opt_weights = tf.keras.backend.batch_get_value( + model.optimizer.weights + ) + else: + model_opt_weights = [var.numpy() for var in model.optimizer.variables()] + if hasattr(loaded_model.optimizer, "weights"): + loaded_opt_weights = tf.keras.backend.batch_get_value( + loaded_model.optimizer.weights + ) + else: + loaded_opt_weights = [ + var.numpy() for var in loaded_model.optimizer.variables() + ] # Assert optimizer weights are equal for weight_model, weight_loaded_model in zip( @@ -277,7 +322,7 @@ def test_save_load_with_sequence_features(self, tmpdir, loss, optimizer, metrics indices_a[:, 0] = np.arange(10) inputs_a = tf.SparseTensor(indices_a, values_a, (batch_size, timesteps, 1)) - values_b = np.zeros(10, dtype=np.str) + values_b = np.zeros(10, dtype=str) indices_b = np.zeros((10, 3), dtype=np.int64) indices_b[:, 0] = np.arange(10) inputs_b = tf.SparseTensor(indices_b, values_b, (batch_size, timesteps, 1)) @@ -310,7 +355,7 @@ def test_functional_model_save_load_with_custom_loss_and_metric(self, tmpdir): tiledb_uri = os.path.join(tmpdir, "model_array") tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri, model=model) tiledb_model_obj.save(include_optimizer=True) - loaded_model = tiledb_model_obj.load(compile_model=True) + loaded_model = tiledb_model_obj.load(compile_model=True, safe_mode=False) # Assert all evaluation results are the same. assert all( diff --git a/tiledb/ml/models/tensorflow_keras.py b/tiledb/ml/models/tensorflow_keras.py index 95c4b73d..06ee7eab 100644 --- a/tiledb/ml/models/tensorflow_keras.py +++ b/tiledb/ml/models/tensorflow_keras.py @@ -15,11 +15,39 @@ from ._base import Meta, TileDBArtifact, Timestamp -FunctionalOrSequential = (keras.models.Functional, keras.models.Sequential) -TFOptimizer = keras.optimizers.TFOptimizer -get_json_type = keras.saving.saved_model.json_utils.get_json_type -preprocess_weights_for_loading = keras.saving.hdf5_format.preprocess_weights_for_loading -saving_utils = keras.saving.saving_utils +keras_major, keras_minor, keras_patch = keras.__version__.split(".") +FunctionalOrSequential = keras.models.Sequential +# Handle keras <=v2.10 +if int(keras_major) <= 2 and int(keras_minor) <= 10: + FunctionalOrSequential = (keras.models.Functional, keras.models.Sequential) + TFOptimizer = keras.optimizers.TFOptimizer + get_json_type = keras.saving.saved_model.json_utils.get_json_type + preprocess_weights_for_loading = ( + keras.saving.hdf5_format.preprocess_weights_for_loading + ) + saving_utils = keras.saving.saving_utils +# Handle keras >=v2.11 +elif int(keras_major) <= 2 and int(keras_minor) <= 12: + FunctionalOrSequential = (keras.models.Functional, keras.models.Sequential) + TFOptimizer = tf.keras.optimizers.legacy.Optimizer + get_json_type = keras.saving.legacy.saved_model.json_utils.get_json_type + preprocess_weights_for_loading = ( + keras.saving.legacy.hdf5_format.preprocess_weights_for_loading + ) + saving_utils = keras.saving.legacy.saving_utils +else: + from keras.src.saving.serialization_lib import SafeModeScope + + FunctionalOrSequential = ( + keras.src.engine.functional.Functional, + keras.src.engine.sequential.Sequential, + ) + TFOptimizer = tf.keras.optimizers.legacy.Optimizer + get_json_type = keras.src.saving.legacy.saved_model.json_utils.get_json_type + preprocess_weights_for_loading = ( + keras.src.saving.legacy.hdf5_format.preprocess_weights_for_loading + ) + saving_utils = keras.src.saving.legacy.saving_utils class TensorflowKerasTileDBModel(TileDBArtifact[tf.keras.Model]): @@ -59,7 +87,7 @@ def save( if not isinstance(self.artifact, FunctionalOrSequential): raise RuntimeError( - "Subclassed Models (Custom Layers) not supported at the moment." + f"Subclassed Models (Custom Layers) for {type(self.artifact)} not supported at the moment." ) # Used in this format only when model is Functional or Sequential @@ -109,6 +137,7 @@ def load( custom_objects: Optional[Mapping[str, Any]] = None, input_shape: Optional[Tuple[int, ...]] = None, callback: bool = False, + safe_mode: Optional[bool] = None, ) -> tf.keras.Model: """ Load switch, i.e, decide between __load (TileDB-ML<=0.8.0) or __load_v2 (TileDB-ML>0.8.0). @@ -129,7 +158,7 @@ def load( model_array, compile_model, callback, custom_objects ) else: - return self.__load(model_array, compile_model, callback) + return self.__load(model_array, compile_model, callback, safe_mode) def __load_legacy( self, @@ -200,13 +229,25 @@ def __load_legacy( return model def __load( - self, model_array: tiledb.Array, compile_model: bool, callback: bool + self, + model_array: tiledb.Array, + compile_model: bool, + callback: bool, + safe_mode: Optional[bool], ) -> tf.keras.Model: model_config = json.loads(model_array.meta["model_config"]) model_class = model_config["class_name"] cls = tf.keras.Sequential if model_class == "Sequential" else tf.keras.Model - model = cls.from_config(model_config["config"]) + + if int(keras_major) <= 2 and int(keras_minor) >= 13: + if safe_mode is not None: + with SafeModeScope(safe_mode=safe_mode): + model = cls.from_config(model_config["config"]) + else: + model = cls.from_config(model_config["config"]) + else: + model = cls.from_config(model_config["config"]) model_weights = self._get_model_param(model_array, "model") model.set_weights(model_weights) @@ -266,6 +307,9 @@ def _serialize_optimizer_weights( assert self.artifact optimizer = self.artifact.optimizer if optimizer and not isinstance(optimizer, TFOptimizer): - optimizer_weights = tf.keras.backend.batch_get_value(optimizer.weights) + if hasattr(optimizer, "weights"): + optimizer_weights = tf.keras.backend.batch_get_value(optimizer.weights) + else: + optimizer_weights = [var.numpy() for var in optimizer.variables()] return pickle.dumps(optimizer_weights, protocol=4) return b""