Merge pull request #399 from bioimage-io/improve_pytorch_adapter

Improve model adapters

Author: FynnBe

bioimageio/core/_resource_tests.py

@@ -1,5 +1,6 @@
 import traceback
 import warnings
+from itertools import product
 from typing import Dict, Hashable, List, Literal, Optional, Set, Tuple, Union
 
 import numpy as np
@@ -179,31 +180,38 @@ def _test_model_inference_parametrized(
     model: v0_5.ModelDescr,
     weight_format: Optional[WeightsFormat],
     devices: Optional[List[str]],
-    test_cases: Set[Tuple[v0_5.ParameterizedSize.N, BatchSize]] = {
-        (0, 2),
-        (1, 3),
-        (2, 1),
-        (3, 2),
-    },
 ) -> None:
-    if not test_cases:
-        return
-
-    logger.info(
-        "Testing inference with {} different input tensor sizes", len(test_cases)
-    )
-
     if not any(
         isinstance(a.size, v0_5.ParameterizedSize)
         for ipt in model.inputs
         for a in ipt.axes
     ):
         # no parameterized sizes => set n=0
-        test_cases = {(0, b) for _n, b in test_cases}
+        ns: Set[v0_5.ParameterizedSize.N] = {0}
+    else:
+        ns = {0, 1, 2}
 
-    if not any(isinstance(a, v0_5.BatchAxis) for ipt in model.inputs for a in ipt.axes):
-        # no batch axis => set b=1
-        test_cases = {(n, 1) for n, _b in test_cases}
+    given_batch_sizes = {
+        a.size
+        for ipt in model.inputs
+        for a in ipt.axes
+        if isinstance(a, v0_5.BatchAxis)
+    }
+    if given_batch_sizes:
+        batch_sizes = {gbs for gbs in given_batch_sizes if gbs is not None}
+        if not batch_sizes:
+            # only arbitrary batch sizes
+            batch_sizes = {1, 2}
+    else:
+        # no batch axis
+        batch_sizes = {1}
+
+    test_cases: Set[Tuple[v0_5.ParameterizedSize.N, BatchSize]] = {
+        (n, b) for n, b in product(sorted(ns), sorted(batch_sizes))
+    }
+    logger.info(
+        "Testing inference with {} different input tensor sizes", len(test_cases)
+    )
 
     def generate_test_cases():
         tested: Set[Hashable] = set()

bioimageio/core/model_adapters/_pytorch_model_adapter.py

@@ -42,7 +42,7 @@ def __init__(
 
         self._primary_device = self._devices[0]
         state: Any = torch.load(
-            download(weights.source).path,
+            download(weights).path,
             map_location=self._primary_device,  # pyright: ignore[reportUnknownArgumentType]
         )
         self._network.load_state_dict(state)

bioimageio/core/model_adapters/_tensorflow_model_adapter.py

@@ -1,8 +1,8 @@
-import warnings
 import zipfile
 from typing import List, Literal, Optional, Sequence, Union
 
 import numpy as np
+from loguru import logger
 
 from bioimageio.spec.common import FileSource
 from bioimageio.spec.model import v0_4, v0_5
@@ -46,19 +46,19 @@ def __init__(
         )
         model_tf_version = weights.tensorflow_version
         if model_tf_version is None:
-            warnings.warn(
+            logger.warning(
                 "The model does not specify the tensorflow version."
                 + f"Cannot check if it is compatible with intalled tensorflow {tf_version}."
             )
         elif model_tf_version > tf_version:
-            warnings.warn(
+            logger.warning(
                 f"The model specifies a newer tensorflow version than installed: {model_tf_version} > {tf_version}."
             )
         elif (model_tf_version.major, model_tf_version.minor) != (
             tf_version.major,
             tf_version.minor,
         ):
-            warnings.warn(
+            logger.warning(
                 "The tensorflow version specified by the model does not match the installed: "
                 + f"{model_tf_version} != {tf_version}."
             )
@@ -70,7 +70,7 @@ def __init__(
 
         # TODO tf device management
         if devices is not None:
-            warnings.warn(
+            logger.warning(
                 f"Device management is not implemented for tensorflow yet, ignoring the devices {devices}"
             )
 
@@ -98,9 +98,20 @@ def _get_network(  # pyright: ignore[reportUnknownParameterType]
         weight_file = self.require_unzipped(weight_file)
         assert tf is not None
         if self.use_keras_api:
-            return tf.keras.models.load_model(
-                weight_file, compile=False
-            )  # pyright: ignore[reportUnknownVariableType]
+            try:
+                return tf.keras.layers.TFSMLayer(
+                    weight_file, call_endpoint="serve"
+                )  # pyright: ignore[reportUnknownVariableType]
+            except Exception as e:
+                try:
+                    return tf.keras.layers.TFSMLayer(
+                        weight_file, call_endpoint="serving_default"
+                    )  # pyright: ignore[reportUnknownVariableType]
+                except Exception as ee:
+                    logger.opt(exception=ee).info(
+                        "keras.layers.TFSMLayer error for alternative call_endpoint='serving_default'"
+                    )
+                    raise e
         else:
             # NOTE in tf1 the model needs to be loaded inside of the session, so we cannot preload the model
             return str(weight_file)
@@ -189,24 +200,15 @@ def _forward_keras(  # pyright: ignore[reportUnknownParameterType]
             None if ipt is None else tf.convert_to_tensor(ipt) for ipt in input_tensors
         ]
 
-        try:
-            result = (  # pyright: ignore[reportUnknownVariableType]
-                self._network.forward(*tf_tensor)
-            )
-        except AttributeError:
-            result = (  # pyright: ignore[reportUnknownVariableType]
-                self._network.predict(*tf_tensor)
-            )
+        result = self._network(*tf_tensor)  # pyright: ignore[reportUnknownVariableType]
 
-        if not isinstance(result, (tuple, list)):
-            result = [result]  # pyright: ignore[reportUnknownVariableType]
+        assert isinstance(result, dict)
+
+        # TODO: Use RDF's `outputs[i].id` here
+        result = list(result.values())
 
         return [  # pyright: ignore[reportUnknownVariableType]
-            (
-                None
-                if r is None
-                else r if isinstance(r, np.ndarray) else tf.make_ndarray(r)
-            )
+            (None if r is None else r if isinstance(r, np.ndarray) else r.numpy())
             for r in result  # pyright: ignore[reportUnknownVariableType]
         ]
 
@@ -230,7 +232,7 @@ def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]:
         ]
 
     def unload(self) -> None:
-        warnings.warn(
+        logger.warning(
             "Device management is not implemented for keras yet, cannot unload model"
         )
 

dev/env-py38.yaml

@@ -4,7 +4,7 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - bioimageio.spec>=0.5.2.post1
+  - bioimageio.spec>=0.5.3
   - black
   - crick # uncommented
   - filelock

dev/env-tf.yaml

@@ -4,7 +4,7 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - bioimageio.spec>=0.5.2.post1
+  - bioimageio.spec>=0.5.3
   - black
   # - crick  # currently requires python<=3.9
   - filelock

dev/env-wo-python.yaml

@@ -4,7 +4,7 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - bioimageio.spec>=0.5.2.post1
+  - bioimageio.spec>=0.5.3
   - black
   # - crick  # currently requires python<=3.9
   - filelock

dev/env.yaml

@@ -3,7 +3,7 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - bioimageio.spec>=0.5.2.post1
+  - bioimageio.spec>=0.5.3
   - black
   # - crick  # currently requires python<=3.9
   - filelock