Reduce gpu memory usage for sleap model...

Author: isaacrobinson2000

RELEASE_NOTES.md

@@ -1,3 +1 @@
-This version of diplomat:
- - Enhanced skeletal transition curve that takes into account spread of body parts.
- - Support for converting DLC hdf5 and sleap hdf5 pose files to diplomat csvs, and passing them directly to the tweak and annotate commands.
+Reduces GPU memory usage when running inference on top of sleap.

diplomat/__init__.py

@@ -2,7 +2,7 @@
 A tool providing multi-animal tracking capabilities on top of other Deep learning based tracking software.
 """
 
-__version__ = "0.3.6"
+__version__ = "0.3.7"
 # Can be used by functions to determine if diplomat was invoked through it's CLI interface.
 CLI_RUN = False
 

diplomat/frontends/deeplabcut/load_model.py

@@ -1,4 +1,3 @@
-import functools
 import tempfile
 from io import BytesIO
 from pathlib import Path
@@ -153,7 +152,6 @@ def _get_dlc_inputs_and_outputs(meta_path):
 
 def _load_and_convert_model(model_dir: Path, device_index: Optional[int], use_cpu: bool):
     import tensorflow as tf
-    from tensorflow.python.training import py_checkpoint_reader
     import tensorflow.compat.v1 as tf_v1
     import tf2onnx
     tf.compat.v1.disable_eager_execution()
@@ -166,7 +164,6 @@ def _load_and_convert_model(model_dir: Path, device_index: Optional[int], use_cp
     latest_meta_file = max(meta_files, key=lambda k: int(k.stem.split("-")[-1]))
 
     inputs, outputs = _get_dlc_inputs_and_outputs(str(latest_meta_file))
-    print(inputs, outputs)
 
     graph_def, inputs, outputs = from_checkpoint(
         str(latest_meta_file), inputs, outputs

diplomat/frontends/sleap/run_utils.py

@@ -1,3 +1,4 @@
+import contextlib
 from .sleap_imports import tf, h5py
 import json
 import zipfile
@@ -127,24 +128,26 @@ def _load_configs_from_zip(z: zipfile.ZipFile, include_model = True):
 
 @resolve_lazy_imports
 def _load_config_and_model(path, include_model = True):
-    path = Path(path)
-    if(zipfile.is_zipfile(path)):
-        with zipfile.ZipFile(path, "r") as z:
-            return _load_configs_from_zip(z, include_model)
-
-    if(path.is_dir()):
-        path = path / "training_config.json"
-    path = path.resolve()
-
-    with path.open("rb") as f:
-        cfg = json.load(f)
-        _correct_skeletons_in_config(cfg)
-    model_path = _resolve_model_path(path.parent.iterdir())
-    if(include_model):
-        model = tf.keras.models.load_model(model_path, compile=False)
-        return [(cfg, model)]
-    else:
-        return [cfg]
+    device_ctx = tf.device('/cpu:0') if include_model else contextlib.nullcontext()
+    with device_ctx:
+        path = Path(path)
+        if(zipfile.is_zipfile(path)):
+            with zipfile.ZipFile(path, "r") as z:
+                return _load_configs_from_zip(z, include_model)
+
+        if(path.is_dir()):
+            path = path / "training_config.json"
+        path = path.resolve()
+
+        with path.open("rb") as f:
+            cfg = json.load(f)
+            _correct_skeletons_in_config(cfg)
+        model_path = _resolve_model_path(path.parent.iterdir())
+        if(include_model):
+            model = tf.keras.models.load_model(model_path, compile=False)
+            return [(cfg, model)]
+        else:
+            return [cfg]
 
 
 def _load_configs(paths, include_models: bool = True):

diplomat/frontends/sleap/sleap_providers.py

@@ -239,31 +239,33 @@ def _reset_input_layer(
     Returns:
         A copy of `keras_model` with input shape `new_shape`.
     """
-
-    if new_shape is None:
-        new_shape = (None, None, None, keras_model.input_shape[-1])
-
-    model_config = keras_model.get_config()
-    model_config["layers"][0]["config"]["batch_input_shape"] = new_shape
-    new_model = tf.keras.Model.from_config(
-        model_config, custom_objects={}
-    )  # Change custom objects if necessary
-
-    # Iterate over all the layers that we want to get weights from
-    weights = [layer.get_weights() for layer in keras_model.layers]
-    for layer, weight in zip(new_model.layers, weights):
-        if len(weight) > 0:
-            layer.set_weights(weight)
+    with tf.device("/cpu:0"):
+        if new_shape is None:
+            new_shape = (None, None, None, keras_model.input_shape[-1])
+
+        model_config = keras_model.get_config()
+        model_config["layers"][0]["config"]["batch_input_shape"] = new_shape
+        new_model = tf.keras.Model.from_config(
+            model_config, custom_objects={}
+        )  # Change custom objects if necessary
+
+        # Iterate over all the layers that we want to get weights from
+        weights = [layer.get_weights() for layer in keras_model.layers]
+        for layer, weight in zip(new_model.layers, weights):
+            if len(weight) > 0:
+                layer.set_weights(weight)
 
     return new_model
 
 
 @resolve_lazy_imports
 def _keras_to_onnx_model(keras_model) -> onnx.ModelProto:
-    input_signature = [
-        tf.TensorSpec(keras_model.input_shape, tf.float32, name="image")
-    ]
-    return tf2onnx.convert.from_keras(keras_model, input_signature, opset=17)[0]
+    with tf.device("/cpu:0"):
+        input_signature = [
+            tf.TensorSpec(keras_model.input_shape, tf.float32, name="image")
+        ]
+        mdl = tf2onnx.convert.from_keras(keras_model, input_signature, opset=17)[0]
+        return mdl
 
 
 def _find_model_output(model: ort.InferenceSession, name: str, required: bool = True):