Merge pull request #155 from bioimage-io/unload

add unload for PredictionPipeline and ModelAdapters

Author: FynnBe

bioimageio/core/prediction_pipeline/_model_adapters/_keras_model_adapter.py

@@ -1,24 +1,26 @@
 import warnings
-from typing import List, Optional
+from typing import List, Optional, Sequence
 
 import keras
 import xarray as xr
 
-from bioimageio.core.resource_io import nodes
 from ._model_adapter import ModelAdapter
 
 
 class KerasModelAdapter(ModelAdapter):
-    def __init__(self, *, bioimageio_model: nodes.Model, devices: Optional[List[str]] = None):
+    def _load(self, *, devices: Optional[Sequence[str]] = None) -> None:
         # TODO keras device management
         if devices is not None:
-            warnings.warn(f"Device management is not implemented for tensorflow yet, ignoring the devices {devices}")
+            warnings.warn(f"Device management is not implemented for keras yet, ignoring the devices {devices}")
 
-        weight_file = bioimageio_model.weights["keras_hdf5"].source
+        weight_file = self.bioimageio_model.weights["keras_hdf5"].source
         self._model = keras.models.load_model(weight_file)
-        self._output_axes = [tuple(out.axes) for out in bioimageio_model.outputs]
+        self._output_axes = [tuple(out.axes) for out in self.bioimageio_model.outputs]
 
-    def forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
+    def _unload(self) -> None:
+        warnings.warn("Device management is not implemented for keras yet, cannot unload model")
+
+    def _forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
         result = self._model.predict(*input_tensors)
         if not isinstance(result, (tuple, list)):
             result = [result]

bioimageio/core/prediction_pipeline/_model_adapters/_model_adapter.py

@@ -1,11 +1,11 @@
 import abc
-from typing import List, Optional, Type
+from typing import List, Optional, Sequence, Type
 
 import xarray as xr
 
 from bioimageio.core.resource_io import nodes
 
-#: Known weigh types in order of priority
+#: Known weight formats in order of priority
 #: First match wins
 _WEIGHT_FORMATS = ["pytorch_state_dict", "tensorflow_saved_model_bundle", "pytorch_script", "onnx", "keras_hdf5"]
 
@@ -15,19 +15,75 @@ class ModelAdapter(abc.ABC):
     Represents model *without* any preprocessing and postprocessing
     """
 
+    def __init__(self, *, bioimageio_model: nodes.Model, devices: Optional[Sequence[str]] = None):
+        self.bioimageio_model = bioimageio_model
+        self.default_devices = devices
+        self.loaded = False
+
+    def __enter__(self):
+        """load on entering context"""
+        assert not self.loaded
+        self.load()  # using default_devices
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """unload on exiting context"""
+        assert self.loaded
+        self.unload()
+        return False
+
+    def load(self, *, devices: Optional[Sequence[str]] = None) -> None:
+        """
+        Note: Use ModelAdapter as context to not worry about calling unload()!
+        Load model onto devices. If devices is None, self.default_devices are chosen
+        (which may be None as well, in which case a framework dependent default is chosen)
+        """
+        self._load(devices=devices or self.default_devices)
+        self.loaded = True
+
     @abc.abstractmethod
-    def __init__(self, *, bioimageio_model: nodes.Model, devices=Optional[List[str]]):
+    def _load(self, *, devices: Optional[Sequence[str]] = None) -> None:
+        """
+        Load model onto devices. If devices is None a framework dependent default is chosen
+        """
         ...
 
-    @abc.abstractmethod
     def forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
+        """
+        Load model if unloaded/outside context; then run forward pass of model to get model predictions
+        """
+        if not self.loaded:
+            self.load()
+
+        assert self.loaded
+        return self._forward(*input_tensors)
+
+    @abc.abstractmethod
+    def _forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
         """
         Run forward pass of model to get model predictions
         Note: model is responsible converting it's data representation to
         xarray.DataArray
         """
         ...
 
+    def unload(self):
+        """
+        Unload model from any devices, freeing their memory.
+        Note: Use ModelAdapter as context to not worry about calling unload()!
+        """
+        # implementation of non-state-machine logic in _unload()
+        assert self.loaded
+        self._unload()
+        self.loaded = False
+
+    @abc.abstractmethod
+    def _unload(self) -> None:
+        """
+        Unload model from any devices, freeing their memory.
+        """
+        ...
+
 
 def get_weight_formats() -> List[str]:
     """
@@ -37,14 +93,13 @@ def get_weight_formats() -> List[str]:
 
 
 def create_model_adapter(
-    *, bioimageio_model: nodes.Model, devices=Optional[List[str]], weight_format: Optional[str] = None
+    *, bioimageio_model: nodes.Model, devices=Optional[Sequence[str]], weight_format: Optional[str] = None
 ) -> ModelAdapter:
     """
     Creates model adapter based on the passed spec
     Note: All specific adapters should happen inside this function to prevent different framework
     initializations interfering with each other
     """
-    spec = bioimageio_model
     weights = bioimageio_model.weights
     weight_formats = get_weight_formats()
 
@@ -59,7 +114,7 @@ def create_model_adapter(
             return adapter_cls(bioimageio_model=bioimageio_model, devices=devices)
 
     raise RuntimeError(
-        f"weight format {weight_format} not among weight formats listed in model: {list(spec.weights.keys())}"
+        f"weight format {weight_format} not among weight formats listed in model: {list(bioimageio_model.weights.keys())}"
     )
 
 

bioimageio/core/prediction_pipeline/_model_adapters/_onnx_model_adapter.py

@@ -5,30 +5,30 @@
 import onnxruntime as rt
 import xarray as xr
 
-from bioimageio.core.resource_io import nodes
 from ._model_adapter import ModelAdapter
 
 logger = logging.getLogger(__name__)
 
 
 class ONNXModelAdapter(ModelAdapter):
-    def __init__(self, *, bioimageio_model: nodes.Model, devices: Optional[List[str]] = None):
-        spec = bioimageio_model
+    def _load(self, *, devices: Optional[List[str]] = None):
+        self._internal_output_axes = [tuple(out.axes) for out in self.bioimageio_model.outputs]
 
-        self._internal_output_axes = [tuple(out.axes) for out in bioimageio_model.outputs]
-
-        self._session = rt.InferenceSession(str(spec.weights["onnx"].source))
+        self._session = rt.InferenceSession(str(self.bioimageio_model.weights["onnx"].source))
         onnx_inputs = self._session.get_inputs()
         self._input_names = [ipt.name for ipt in onnx_inputs]
 
         if devices is not None:
             warnings.warn(f"Device management is not implemented for onnx yet, ignoring the devices {devices}")
 
-    def forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
+    def _forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
         assert len(input_tensors) == len(self._input_names)
         input_arrays = [ipt.data for ipt in input_tensors]
         result = self._session.run(None, dict(zip(self._input_names, input_arrays)))
         if not isinstance(result, (list, tuple)):
             result = []
 
         return [xr.DataArray(r, dims=axes) for r, axes in zip(result, self._internal_output_axes)]
+
+    def _unload(self) -> None:
+        warnings.warn("Device management is not implemented for onnx yet, cannot unload model")

bioimageio/core/prediction_pipeline/_model_adapters/_pytorch_model_adapter.py

@@ -1,3 +1,5 @@
+import gc
+import warnings
 from typing import List, Optional
 
 import torch
@@ -9,24 +11,28 @@
 
 
 class PytorchModelAdapter(ModelAdapter):
-    def __init__(self, *, bioimageio_model: nodes.Model, devices: Optional[List[str]] = None):
-        self._model = self.get_nn_instance(bioimageio_model)
+    def _load(self, *, devices: Optional[List[str]] = None):
+        self._model = self.get_nn_instance(self.bioimageio_model)
 
         if devices is None:
             self._devices = ["cuda" if torch.cuda.is_available() else "cpu"]
         else:
             self._devices = [torch.device(d) for d in devices]
+
+        if len(self._devices) > 1:
+            warnings.warn("Multiple devices for single pytorch model not yet implemented")
+
         self._model.to(self._devices[0])
 
         assert isinstance(self._model, torch.nn.Module)
-        weights = bioimageio_model.weights.get("pytorch_state_dict")
+        weights = self.bioimageio_model.weights.get("pytorch_state_dict")
         if weights is not None and weights.source:
             state = torch.load(weights.source, map_location=self._devices[0])
             self._model.load_state_dict(state)
 
-        self._internal_output_axes = [tuple(out.axes) for out in bioimageio_model.outputs]
+        self._internal_output_axes = [tuple(out.axes) for out in self.bioimageio_model.outputs]
 
-    def forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
+    def _forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
         with torch.no_grad():
             tensors = [torch.from_numpy(ipt.data) for ipt in input_tensors]
             tensors = [t.to(self._devices[0]) for t in tensors]
@@ -38,6 +44,12 @@ def forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
 
         return [xr.DataArray(r, dims=axes) for r, axes in zip(result, self._internal_output_axes)]
 
+    def _unload(self) -> None:
+        self._devices = None
+        del self._model
+        gc.collect()  # deallocate memory
+        torch.cuda.empty_cache()  # release reserved memory
+
     @staticmethod
     def get_nn_instance(model_node: nodes.Model, **kwargs):
         assert isinstance(model_node.source, nodes.ImportedSource)

bioimageio/core/prediction_pipeline/_model_adapters/_tensorflow_model_adapter.py

@@ -9,8 +9,15 @@
 from bioimageio.core.resource_io import nodes
 from ._model_adapter import ModelAdapter
 
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal  # type: ignore
+
 
 class TensorflowModelAdapterBase(ModelAdapter):
+    weight_format: Literal["keras_hdf5", "tensorflow_saved_model_bundle"]
+
     def require_unzipped(self, weight_file):
         if zipfile.is_zipfile(weight_file):
             out_path = weight_file.with_suffix("")
@@ -27,30 +34,28 @@ def _load_model(self, weight_file):
             # NOTE in tf1 the model needs to be loaded inside of the session, so we cannot preload the model
             return str(weight_file)
 
-    def __init__(self, *, bioimageio_model: nodes.Model, weight_format: str, devices: Optional[List[str]] = None):
-        self.spec = bioimageio_model
-
+    def _load(self, *, devices: Optional[List[str]] = None):
         try:
-            tf_version = self.spec.weights[weight_format].tensorflow_version.version
+            tf_version = self.bioimageio_model.weights[self.weight_format].tensorflow_version.version
         except AttributeError:
             tf_version = (1, 14, 0)
         tf_major_ver = tf_version[0]
         assert tf_major_ver in (1, 2)
-        self.use_keras_api = tf_major_ver > 1 or weight_format == "keras_hdf5"
+        self.use_keras_api = tf_major_ver > 1 or self.weight_format == KerasModelAdapter.weight_format
 
         # TODO tf device management
         if devices is not None:
             warnings.warn(f"Device management is not implemented for tensorflow yet, ignoring the devices {devices}")
 
-        weight_file = self.require_unzipped(self.spec.weights[weight_format].source)
+        weight_file = self.require_unzipped(self.bioimageio_model.weights[self.weight_format].source)
         self._model = self._load_model(weight_file)
-        self._internal_output_axes = [tuple(out.axes) for out in bioimageio_model.outputs]
+        self._internal_output_axes = [tuple(out.axes) for out in self.bioimageio_model.outputs]
 
     # TODO currently we relaod the model every time. it would be better to keep the graph and session
     # alive in between of forward passes (but then the sessions need to be properly opened / closed)
     def _forward_tf(self, *input_tensors):
-        input_keys = [ipt.name for ipt in self.spec.inputs]
-        output_keys = [out.name for out in self.spec.outputs]
+        input_keys = [ipt.name for ipt in self.bioimageio_model.inputs]
+        output_keys = [out.name for out in self.bioimageio_model.outputs]
 
         # TODO read from spec
         tag = tf.saved_model.tag_constants.SERVING
@@ -85,7 +90,7 @@ def _forward_keras(self, input_tensors):
 
         return [r if isinstance(r, np.ndarray) else tf.make_ndarray(r) for r in result]
 
-    def forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
+    def _forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
         data = [ipt.data for ipt in input_tensors]
         if self.use_keras_api:
             result = self._forward_keras(*data)
@@ -94,14 +99,13 @@ def forward(self, *input_tensors: xr.DataArray) -> List[xr.DataArray]:
 
         return [xr.DataArray(r, dims=axes) for r, axes in zip(result, self._internal_output_axes)]
 
+    def _unload(self) -> None:
+        warnings.warn("Device management is not implemented for keras yet, cannot unload model")
+
 
 class TensorflowModelAdapter(TensorflowModelAdapterBase):
-    def __init__(self, *, bioimageio_model: nodes.Model, devices=List[str]):
-        weight_format = "tensorflow_saved_model_bundle"
-        super().__init__(bioimageio_model=bioimageio_model, weight_format=weight_format, devices=devices)
+    weight_format = "tensorflow_saved_model_bundle"
 
 
 class KerasModelAdapter(TensorflowModelAdapterBase):
-    def __init__(self, *, bioimageio_model: nodes.Model, devices=List[str]):
-        weight_format = "keras_hdf5"
-        super().__init__(bioimageio_model=bioimageio_model, weight_format=weight_format, devices=devices)
+    weight_format = "keras_hdf5"

bioimageio/core/prediction_pipeline/_model_adapters/_torchscript_model_adapter.py

@@ -1,26 +1,30 @@
+import gc
+import warnings
 from typing import List, Optional
 
 import numpy as np
 import torch
 import xarray as xr
 
-from bioimageio.core.resource_io import nodes
 from ._model_adapter import ModelAdapter
 
 
 class TorchscriptModelAdapter(ModelAdapter):
-    def __init__(self, *, bioimageio_model: nodes.Model, devices: Optional[List[str]] = None):
-        weight_path = str(bioimageio_model.weights["pytorch_script"].source.resolve())
+    def _load(self, *, devices: Optional[List[str]] = None):
+        weight_path = str(self.bioimageio_model.weights["pytorch_script"].source.resolve())
         if devices is None:
             self.devices = ["cuda" if torch.cuda.is_available() else "cpu"]
         else:
             self.devices = [torch.device(d) for d in devices]
 
+        if len(self.devices) > 1:
+            warnings.warn("Multiple devices for single torchscript model not yet implemented")
+
         self._model = torch.jit.load(weight_path)
         self._model.to(self.devices[0])
-        self._internal_output_axes = [tuple(out.axes) for out in bioimageio_model.outputs]
+        self._internal_output_axes = [tuple(out.axes) for out in self.bioimageio_model.outputs]
 
-    def forward(self, *batch: xr.DataArray) -> List[xr.DataArray]:
+    def _forward(self, *batch: xr.DataArray) -> List[xr.DataArray]:
         with torch.no_grad():
             torch_tensor = [torch.from_numpy(b.data).to(self.devices[0]) for b in batch]
             result = self._model.forward(*torch_tensor)
@@ -31,3 +35,9 @@ def forward(self, *batch: xr.DataArray) -> List[xr.DataArray]:
 
         assert len(result) == len(self._internal_output_axes)
         return [xr.DataArray(r, dims=axes) for r, axes in zip(result, self._internal_output_axes)]
+
+    def _unload(self) -> None:
+        self._devices = None
+        del self._model
+        gc.collect()  # deallocate memory
+        torch.cuda.empty_cache()  # release reserved memory