Fix weight converters and return their corresponding v5 weight descr

Author: thodkatz

bioimageio/core/weight_converter/keras/_tensorflow.py

@@ -5,6 +5,9 @@
 from typing import no_type_check
 from zipfile import ZipFile
 
+from bioimageio.spec._internal.version_type import Version
+from bioimageio.spec.model import v0_5
+
 try:
     import tensorflow.saved_model
 except Exception:
@@ -39,7 +42,7 @@ def _convert_tf1(
     input_name: str,
     output_name: str,
     zip_weights: bool,
-):
+) -> v0_5.TensorflowSavedModelBundleWeightsDescr:
     try:
         # try to build the tf model with the keras import from tensorflow
         from bioimageio.core.weight_converter.keras._tensorflow import (
@@ -77,10 +80,16 @@ def build_tf_model():
         output_path = _zip_model_bundle(output_path)
     print("TensorFlow model exported to", output_path)
 
-    return 0
+    return v0_5.TensorflowSavedModelBundleWeightsDescr(
+        source=output_path,
+        parent="keras_hdf5",
+        tensorflow_version=Version(tensorflow.__version__),
+    )
 
 
-def _convert_tf2(keras_weight_path: Path, output_path: Path, zip_weights: bool):
+def _convert_tf2(
+    keras_weight_path: Path, output_path: Path, zip_weights: bool
+) -> v0_5.TensorflowSavedModelBundleWeightsDescr:
     try:
         # try to build the tf model with the keras import from tensorflow
         from bioimageio.core.weight_converter.keras._tensorflow import keras
@@ -95,12 +104,16 @@ def _convert_tf2(keras_weight_path: Path, output_path: Path, zip_weights: bool):
         output_path = _zip_model_bundle(output_path)
     print("TensorFlow model exported to", output_path)
 
-    return 0
+    return v0_5.TensorflowSavedModelBundleWeightsDescr(
+        source=output_path,
+        parent="keras_hdf5",
+        tensorflow_version=tensorflow.__version__,
+    )
 
 
 def convert_weights_to_tensorflow_saved_model_bundle(
     model: ModelDescr, output_path: Path
-):
+) -> v0_5.TensorflowSavedModelBundleWeightsDescr:
     """Convert model weights from format 'keras_hdf5' to 'tensorflow_saved_model_bundle'.
 
     Adapted from

bioimageio/core/weight_converter/torch/_onnx.py

@@ -1,13 +1,11 @@
 # type: ignore  # TODO: type
-import warnings
+from __future__ import annotations
 from pathlib import Path
-from typing import Any, List, Sequence, cast
+from typing import Any, List, Sequence, cast, Union
 
 import numpy as np
 from numpy.testing import assert_array_almost_equal
 
-from bioimageio.spec import load_description
-from bioimageio.spec.common import InvalidDescr
 from bioimageio.spec.model import v0_4, v0_5
 
 from ...digest_spec import get_member_id, get_test_inputs
@@ -19,15 +17,15 @@
     torch = None
 
 
-def add_onnx_weights(
-    model_spec: "str | Path | v0_4.ModelDescr | v0_5.ModelDescr",
+def convert_weights_to_onnx(
+    model_spec: Union[v0_4.ModelDescr, v0_5.ModelDescr],
     *,
     output_path: Path,
     use_tracing: bool = True,
     test_decimal: int = 4,
     verbose: bool = False,
-    opset_version: "int | None" = None,
-):
+    opset_version: int = 15,
+) -> v0_5.OnnxWeightsDescr:
     """Convert model weights from format 'pytorch_state_dict' to 'onnx'.
 
     Args:
@@ -36,16 +34,6 @@ def add_onnx_weights(
         use_tracing: whether to use tracing or scripting to export the onnx format
         test_decimal: precision for testing whether the results agree
     """
-    if isinstance(model_spec, (str, Path)):
-        loaded_spec = load_description(Path(model_spec))
-        if isinstance(loaded_spec, InvalidDescr):
-            raise ValueError(f"Bad resource description: {loaded_spec}")
-        if not isinstance(loaded_spec, (v0_4.ModelDescr, v0_5.ModelDescr)):
-            raise TypeError(
-                f"Path {model_spec} is a {loaded_spec.__class__.__name__}, expected a v0_4.ModelDescr or v0_5.ModelDescr"
-            )
-        model_spec = loaded_spec
-
     state_dict_weights_descr = model_spec.weights.pytorch_state_dict
     if state_dict_weights_descr is None:
         raise ValueError(
@@ -54,9 +42,10 @@ def add_onnx_weights(
 
     assert torch is not None
     with torch.no_grad():
-
         sample = get_test_inputs(model_spec)
-        input_data = [sample[get_member_id(ipt)].data.data for ipt in model_spec.inputs]
+        input_data = [
+            sample.members[get_member_id(ipt)].data.data for ipt in model_spec.inputs
+        ]
         input_tensors = [torch.from_numpy(ipt) for ipt in input_data]
         model = load_torch_model(state_dict_weights_descr)
 
@@ -81,9 +70,9 @@ def add_onnx_weights(
     try:
         import onnxruntime as rt  # pyright: ignore [reportMissingTypeStubs]
     except ImportError:
-        msg = "The onnx weights were exported, but onnx rt is not available and weights cannot be checked."
-        warnings.warn(msg)
-        return
+        raise ImportError(
+            "The onnx weights were exported, but onnx rt is not available and weights cannot be checked."
+        )
 
     # check the onnx model
     sess = rt.InferenceSession(str(output_path))
@@ -101,8 +90,11 @@ def add_onnx_weights(
     try:
         for exp, out in zip(expected_outputs, outputs):
             assert_array_almost_equal(exp, out, decimal=test_decimal)
-        return 0
     except AssertionError as e:
-        msg = f"The onnx weights were exported, but results before and after conversion do not agree:\n {str(e)}"
-        warnings.warn(msg)
-        return 1
+        raise ValueError(
+            f"Results before and after weights conversion do not agree:\n {str(e)}"
+        )
+
+    return v0_5.OnnxWeightsDescr(
+        source=output_path, parent="pytorch_state_dict", opset_version=opset_version
+    )

bioimageio/core/weight_converter/torch/_torchscript.py

@@ -1,9 +1,11 @@
 # type: ignore  # TODO: type
+from __future__ import annotations
 from pathlib import Path
 from typing import List, Sequence, Union
 
 import numpy as np
 from numpy.testing import assert_array_almost_equal
+from torch.jit import ScriptModule
 from typing_extensions import Any, assert_never
 
 from bioimageio.spec.model import v0_4, v0_5
@@ -17,12 +19,11 @@
     torch = None
 
 
-# FIXME: remove Any
 def _check_predictions(
     model: Any,
     scripted_model: Any,
-    model_spec: "v0_4.ModelDescr | v0_5.ModelDescr",
-    input_data: Sequence["torch.Tensor"],
+    model_spec: v0_4.ModelDescr | v0_5.ModelDescr,
+    input_data: Sequence[torch.Tensor],
 ):
     assert torch is not None
 
@@ -77,22 +78,27 @@ def _check(input_: Sequence[torch.Tensor]) -> None:
             else:
                 assert_never(axis.size)
 
-    half_step = [st // 2 for st in step]
+    input_data = input_data[0]
+    max_shape = input_data.shape
     max_steps = 4
 
     # check that input and output agree for decreasing input sizes
     for step_factor in range(1, max_steps + 1):
         slice_ = tuple(
-            slice(None) if st == 0 else slice(step_factor * st, -step_factor * st)
-            for st in half_step
-        )
-        this_input = [inp[slice_] for inp in input_data]
-        this_shape = this_input[0].shape
-        if any(tsh < msh for tsh, msh in zip(this_shape, min_shape)):
-            raise ValueError(
-                f"Mismatched shapes: {this_shape}. Expected at least {min_shape}"
+            (
+                slice(None)
+                if step_dim == 0
+                else slice(0, max_dim - step_factor * step_dim, 1)
             )
-        _check(this_input)
+            for max_dim, step_dim in zip(max_shape, step)
+        )
+        sliced_input = input_data[slice_]
+        if any(
+            sliced_dim < min_dim
+            for sliced_dim, min_dim in zip(sliced_input.shape, min_shape)
+        ):
+            return
+        _check([sliced_input])
 
 
 def convert_weights_to_torchscript(
@@ -107,7 +113,6 @@ def convert_weights_to_torchscript(
         output_path: where to save the torchscript weights
         use_tracing: whether to use tracing or scripting to export the torchscript format
     """
-
     state_dict_weights_descr = model_descr.weights.pytorch_state_dict
     if state_dict_weights_descr is None:
         raise ValueError(
@@ -118,26 +123,20 @@ def convert_weights_to_torchscript(
 
     with torch.no_grad():
         input_data = [torch.from_numpy(inp.astype("float32")) for inp in input_data]
-
         model = load_torch_model(state_dict_weights_descr)
-
-        # FIXME: remove Any
-        if use_tracing:
-            scripted_model: Any = torch.jit.trace(model, input_data)
-        else:
-            scripted_model: Any = torch.jit.script(model)
-
+        scripted_module: ScriptModule = (
+            torch.jit.trace(model, input_data)
+            if use_tracing
+            else torch.jit.script(model)
+        )
         _check_predictions(
             model=model,
-            scripted_model=scripted_model,
+            scripted_model=scripted_module,
             model_spec=model_descr,
             input_data=input_data,
         )
 
-    # save the torchscript model
-    scripted_model.save(
-        str(output_path)
-    )  # does not support Path, so need to cast to str
+    scripted_module.save(str(output_path))
 
     return v0_5.TorchscriptWeightsDescr(
         source=output_path,

setup.py

@@ -47,7 +47,7 @@
     extras_require={
         "pytorch": ["torch>=1.6", "torchvision", "keras>=3.0"],
         "tensorflow": ["tensorflow", "keras>=2.15"],
-        "onnx": ["onnxruntime"],
+        "onnx": ["onnxruntime", "onnx"],
         "dev": [
             "black",
             # "crick",  # currently requires python<=3.9

tests/weight_converter/keras/test_tensorflow.py

@@ -3,49 +3,33 @@
 from pathlib import Path
 
 import pytest
-
 from bioimageio.spec import load_description
-from bioimageio.spec.model.v0_5 import ModelDescr
+from bioimageio.spec.model import v0_5
 
+from bioimageio.core.weight_converter.keras._tensorflow import (
+    convert_weights_to_tensorflow_saved_model_bundle,
+)
 
-@pytest.mark.skip(
-    "tensorflow converter not updated yet"
-)  # TODO: test tensorflow converter
-def test_tensorflow_converter(any_keras_model: Path, tmp_path: Path):
-    from bioimageio.core.weight_converter.keras import (
-        convert_weights_to_tensorflow_saved_model_bundle,
-    )
 
-    out_path = tmp_path / "weights"
+@pytest.mark.skip()
+def test_tensorflow_converter(any_keras_model: Path, tmp_path: Path):
     model = load_description(any_keras_model)
-    assert isinstance(model, ModelDescr), model.validation_summary.format()
+    out_path = tmp_path / "weights.h5"
     ret_val = convert_weights_to_tensorflow_saved_model_bundle(model, out_path)
     assert out_path.exists()
-    assert (out_path / "variables").exists()
-    assert (out_path / "saved_model.pb").exists()
-    assert (
-        ret_val == 0
-    )  # check for correctness is done in converter and returns 0 if it passes
+    assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr)
+    assert ret_val.source == out_path
 
 
-@pytest.mark.skip(
-    "tensorflow converter not updated yet"
-)  # TODO: test tensorflow converter
+@pytest.mark.skip()
 def test_tensorflow_converter_zipped(any_keras_model: Path, tmp_path: Path):
-    from bioimageio.core.weight_converter.keras import (
-        convert_weights_to_tensorflow_saved_model_bundle,
-    )
-
     out_path = tmp_path / "weights.zip"
     model = load_description(any_keras_model)
-    assert isinstance(model, ModelDescr), model.validation_summary.format()
     ret_val = convert_weights_to_tensorflow_saved_model_bundle(model, out_path)
+
     assert out_path.exists()
-    assert (
-        ret_val == 0
-    )  # check for correctness is done in converter and returns 0 if it passes
+    assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr)
 
-    # make sure that the zip package was created correctly
     expected_names = {"saved_model.pb", "variables/variables.index"}
     with zipfile.ZipFile(out_path, "r") as f:
         names = set([name for name in f.namelist()])

tests/weight_converter/torch/test_onnx.py

@@ -1,18 +1,23 @@
 # type: ignore  # TODO enable type checking
 import os
-from pathlib import Path
 
-import pytest
+from bioimageio.spec import load_description
+from bioimageio.spec.model import v0_5
 
+from bioimageio.core.weight_converter.torch._onnx import convert_weights_to_onnx
 
-@pytest.mark.skip("onnx converter not updated yet")  # TODO: test onnx converter
-def test_onnx_converter(convert_to_onnx: Path, tmp_path: Path):
-    from bioimageio.core.weight_converter.torch._onnx import convert_weights_to_onnx
 
+def test_onnx_converter(convert_to_onnx, tmp_path):
+    bio_model = load_description(convert_to_onnx)
     out_path = tmp_path / "weights.onnx"
-    ret_val = convert_weights_to_onnx(convert_to_onnx, out_path, test_decimal=3)
+    opset_version = 15
+    ret_val = convert_weights_to_onnx(
+        model_spec=bio_model,
+        output_path=out_path,
+        test_decimal=3,
+        opset_version=opset_version,
+    )
     assert os.path.exists(out_path)
-    if not pytest.skip_onnx:
-        assert (
-            ret_val == 0
-        )  # check for correctness is done in converter and returns 0 if it passes
+    assert isinstance(ret_val, v0_5.OnnxWeightsDescr)
+    assert ret_val.opset_version == opset_version
+    assert ret_val.source == out_path

tests/weight_converter/torch/test_torchscript.py

@@ -1,22 +1,18 @@
 # type: ignore  # TODO enable type checking
-from pathlib import Path
-
 import pytest
+from bioimageio.spec import load_description
+from bioimageio.spec.model import v0_5
 
-from bioimageio.spec.model import v0_4, v0_5
-
+from bioimageio.core.weight_converter.torch._torchscript import (
+    convert_weights_to_torchscript,
+)
 
-@pytest.mark.skip(
-    "torchscript converter not updated yet"
-)  # TODO: test torchscript converter
-def test_torchscript_converter(
-    any_torch_model: "v0_4.ModelDescr | v0_5.ModelDescr", tmp_path: Path
-):
-    from bioimageio.core.weight_converter.torch import convert_weights_to_torchscript
 
+@pytest.mark.skip()
+def test_torchscript_converter(any_torch_model, tmp_path):
+    bio_model = load_description(any_torch_model)
     out_path = tmp_path / "weights.pt"
-    ret_val = convert_weights_to_torchscript(any_torch_model, out_path)
+    ret_val = convert_weights_to_torchscript(bio_model, out_path)
     assert out_path.exists()
-    assert (
-        ret_val == 0
-    )  # check for correctness is done in converter and returns 0 if it passes
+    assert isinstance(ret_val, v0_5.TorchscriptWeightsDescr)
+    assert ret_val.source == out_path