Updates for Tensorflow 2.9+ & Pytorch 1.13 (#195)

* Support tensorflow>=2.9

* Support pytorch 1.13

* [CI] Test torch 1.13 & tensorflow 2.9/2.10

Author: gsakkis

.github/workflows/ci.yml

@@ -10,17 +10,18 @@ jobs:
       fail-fast: false
       matrix:
         ml-deps:
-          - "torch==1.11.0+cpu torchvision==0.12.0+cpu torchdata==0.3.0 tensorflow-cpu==2.7.1"
-          - "torch==1.12.1+cpu torchvision==0.13.1+cpu torchdata==0.4.1 tensorflow-cpu==2.8.1"
+          - "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"
 
     env:
       run_coverage: ${{ github.ref == 'refs/heads/master' }}
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
 
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: "3.7"
 

tests/models/test_tensorflow_keras_models.py

@@ -12,16 +12,15 @@
 import tensorflow as tf
 
 import tiledb
-from tiledb.ml.models.tensorflow_keras import (
-    TensorflowKerasTileDBModel,
-    tf_keras_is_keras,
-)
-
-if tf_keras_is_keras:
-    from keras import testing_utils
-else:
-    from tensorflow.python.keras import testing_utils
+from tiledb.ml.models.tensorflow_keras import TensorflowKerasTileDBModel
 
+try:
+    from keras.testing_infra.test_utils import (
+        get_small_functional_mlp,
+        get_small_sequential_mlp,
+    )
+except ImportError:
+    from keras.testing_utils import get_small_functional_mlp, get_small_sequential_mlp
 
 # Suppress all Tensorflow messages
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -75,11 +74,7 @@ def test_load_tiledb_error_with_wrong_uri():
 
 api = pytest.mark.parametrize(
     "api",
-    [
-        testing_utils.get_small_sequential_mlp,
-        testing_utils.get_small_functional_mlp,
-        ConfigSubclassModel,
-    ],
+    [get_small_sequential_mlp, get_small_functional_mlp, ConfigSubclassModel],
 )
 
 

tiledb/ml/models/tensorflow_keras.py

@@ -8,6 +8,7 @@
 from operator import attrgetter
 from typing import Any, List, Mapping, Optional, Tuple
 
+import keras
 import numpy as np
 import tensorflow as tf
 
@@ -16,20 +17,9 @@
 from ._base import Meta, TileDBArtifact, Timestamp, current_milli_time, group_create
 from ._tensorboard import TensorBoardTileDB
 
-try:
-    import keras
-
-    if keras.Model is not tf.keras.Model:
-        raise ImportError
-    tf_keras_is_keras = True
-except ImportError:
-    import tensorflow.python.keras as keras
-
-    tf_keras_is_keras = False
-
 SharedObjectLoadingScope = keras.utils.generic_utils.SharedObjectLoadingScope
 FunctionalOrSequential = (keras.models.Functional, keras.models.Sequential)
-TFOptimizer = keras.optimizer_v1.TFOptimizer
+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

tiledb/ml/readers/_pytorch_collators.py

@@ -12,6 +12,12 @@
 from ._tensor_schema import TensorSchema
 from .types import TensorKind
 
+try:
+    nested_tensor = torch.nested.nested_tensor
+except AttributeError:
+    nested_tensor = getattr(torch, "nested_tensor", None)
+
+
 T = TypeVar("T")
 
 
@@ -101,7 +107,7 @@ def convert(self, value: np.ndarray) -> torch.Tensor:
 
     def collate(self, batch: Sequence[np.ndarray]) -> torch.Tensor:
         if self.to_nested:
-            return torch.nested_tensor(tuple(map(torch.from_numpy, batch)))
+            return nested_tensor(list(map(torch.from_numpy, batch)))
         else:
             return torch.from_numpy(np.stack(batch))