Merge pull request #210 from TileDB-Inc/sethshelnutt/sc-31839/tensorflowkerastiledbmodel-doesn-t-load

Update to support tensorflow >= 2.11

Author: Shelnutt2

.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

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(

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""