From fbbfc573e2cd2449a9372c4b5bfe131a2dfc8ee1 Mon Sep 17 00:00:00 2001 From: Theodoros Katzalis Date: Thu, 15 Aug 2024 15:17:45 +0200 Subject: [PATCH 01/47] Fix weight converters and return their corresponding v5 weight descr --- .../weight_converter/keras/_tensorflow.py | 23 ++++++-- .../core/weight_converter/torch/_onnx.py | 46 +++++++--------- .../weight_converter/torch/_torchscript.py | 53 +++++++++---------- setup.py | 2 +- .../weight_converter/keras/test_tensorflow.py | 40 +++++--------- tests/weight_converter/torch/test_onnx.py | 25 +++++---- .../torch/test_torchscript.py | 26 ++++----- 7 files changed, 102 insertions(+), 113 deletions(-) diff --git a/bioimageio/core/weight_converter/keras/_tensorflow.py b/bioimageio/core/weight_converter/keras/_tensorflow.py index c901f458..5804226d 100644 --- a/bioimageio/core/weight_converter/keras/_tensorflow.py +++ b/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 diff --git a/bioimageio/core/weight_converter/torch/_onnx.py b/bioimageio/core/weight_converter/torch/_onnx.py index 3935e1d1..d3c7bf01 100644 --- a/bioimageio/core/weight_converter/torch/_onnx.py +++ b/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 + ) diff --git a/bioimageio/core/weight_converter/torch/_torchscript.py b/bioimageio/core/weight_converter/torch/_torchscript.py index 5ca16069..16dc0128 100644 --- a/bioimageio/core/weight_converter/torch/_torchscript.py +++ b/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, diff --git a/setup.py b/setup.py index 7aa66e16..a1a86f45 100644 --- a/setup.py +++ b/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 diff --git a/tests/weight_converter/keras/test_tensorflow.py b/tests/weight_converter/keras/test_tensorflow.py index 65c93f60..18a4f2dc 100644 --- a/tests/weight_converter/keras/test_tensorflow.py +++ b/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()]) diff --git a/tests/weight_converter/torch/test_onnx.py b/tests/weight_converter/torch/test_onnx.py index 54f2cdf4..faab39d6 100644 --- a/tests/weight_converter/torch/test_onnx.py +++ b/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 diff --git a/tests/weight_converter/torch/test_torchscript.py b/tests/weight_converter/torch/test_torchscript.py index e0cee3d8..6b397f08 100644 --- a/tests/weight_converter/torch/test_torchscript.py +++ b/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 From a37b56879def4802e27b026e58126929cbee2ca2 Mon Sep 17 00:00:00 2001 From: Theodoros Katzalis Date: Thu, 15 Aug 2024 16:03:06 +0200 Subject: [PATCH 02/47] Create an interface for weight conversion - Instead of having one module per conversion, use one module, and individual classes will handle the logic of the conversion - Create an abstract class to have a common interface among conversions --- bioimageio/core/weight_converter/__init__.py | 1 - .../core/weight_converter/keras/__init__.py | 1 - .../weight_converter/keras/_tensorflow.py | 164 ------ .../core/weight_converter/torch/__init__.py | 1 - .../core/weight_converter/torch/_onnx.py | 100 ---- .../weight_converter/torch/_torchscript.py | 145 ------ .../core/weight_converter/torch/_utils.py | 24 - bioimageio/core/weight_converters.py | 492 ++++++++++++++++++ .../test_add_weights.py | 0 tests/test_weight_converters.py | 69 +++ .../weight_converter/keras/test_tensorflow.py | 36 -- tests/weight_converter/torch/test_onnx.py | 23 - .../torch/test_torchscript.py | 18 - 13 files changed, 561 insertions(+), 513 deletions(-) delete mode 100644 bioimageio/core/weight_converter/__init__.py delete mode 100644 bioimageio/core/weight_converter/keras/__init__.py delete mode 100644 bioimageio/core/weight_converter/keras/_tensorflow.py delete mode 100644 bioimageio/core/weight_converter/torch/__init__.py delete mode 100644 bioimageio/core/weight_converter/torch/_onnx.py delete mode 100644 bioimageio/core/weight_converter/torch/_torchscript.py delete mode 100644 bioimageio/core/weight_converter/torch/_utils.py create mode 100644 bioimageio/core/weight_converters.py rename tests/{weight_converter => }/test_add_weights.py (100%) create mode 100644 tests/test_weight_converters.py delete mode 100644 tests/weight_converter/keras/test_tensorflow.py delete mode 100644 tests/weight_converter/torch/test_onnx.py delete mode 100644 tests/weight_converter/torch/test_torchscript.py diff --git a/bioimageio/core/weight_converter/__init__.py b/bioimageio/core/weight_converter/__init__.py deleted file mode 100644 index 5f1674c9..00000000 --- a/bioimageio/core/weight_converter/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""coming soon""" diff --git a/bioimageio/core/weight_converter/keras/__init__.py b/bioimageio/core/weight_converter/keras/__init__.py deleted file mode 100644 index 195b42b8..00000000 --- a/bioimageio/core/weight_converter/keras/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# TODO: update keras weight converters diff --git a/bioimageio/core/weight_converter/keras/_tensorflow.py b/bioimageio/core/weight_converter/keras/_tensorflow.py deleted file mode 100644 index 5804226d..00000000 --- a/bioimageio/core/weight_converter/keras/_tensorflow.py +++ /dev/null @@ -1,164 +0,0 @@ -# type: ignore # TODO: type -import os -import shutil -from pathlib import Path -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: - tensorflow = None - -from bioimageio.spec._internal.io_utils import download -from bioimageio.spec.model.v0_5 import ModelDescr - - -def _zip_model_bundle(model_bundle_folder: Path): - zipped_model_bundle = model_bundle_folder.with_suffix(".zip") - - with ZipFile(zipped_model_bundle, "w") as zip_obj: - for root, _, files in os.walk(model_bundle_folder): - for filename in files: - src = os.path.join(root, filename) - zip_obj.write(src, os.path.relpath(src, model_bundle_folder)) - - try: - shutil.rmtree(model_bundle_folder) - except Exception: - print("TensorFlow bundled model was not removed after compression") - - return zipped_model_bundle - - -# adapted from -# https://github.com/deepimagej/pydeepimagej/blob/master/pydeepimagej/yaml/create_config.py#L236 -def _convert_tf1( - keras_weight_path: Path, - output_path: Path, - 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 ( - keras, # type: ignore - ) - - except Exception: - # if the above fails try to export with the standalone keras - import keras - - @no_type_check - def build_tf_model(): - keras_model = keras.models.load_model(keras_weight_path) - assert tensorflow is not None - builder = tensorflow.saved_model.builder.SavedModelBuilder(output_path) - signature = tensorflow.saved_model.signature_def_utils.predict_signature_def( - inputs={input_name: keras_model.input}, - outputs={output_name: keras_model.output}, - ) - - signature_def_map = { - tensorflow.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature - } - - builder.add_meta_graph_and_variables( - keras.backend.get_session(), - [tensorflow.saved_model.tag_constants.SERVING], - signature_def_map=signature_def_map, - ) - builder.save() - - build_tf_model() - - if zip_weights: - output_path = _zip_model_bundle(output_path) - print("TensorFlow model exported to", output_path) - - 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 -) -> 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 - except Exception: - # if the above fails try to export with the standalone keras - import keras - - model = keras.models.load_model(keras_weight_path) - keras.models.save_model(model, output_path) - - if zip_weights: - output_path = _zip_model_bundle(output_path) - print("TensorFlow model exported to", output_path) - - 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 - https://github.com/deepimagej/pydeepimagej/blob/5aaf0e71f9b04df591d5ca596f0af633a7e024f5/pydeepimagej/yaml/create_config.py - - Args: - model: The bioimageio model description - output_path: where to save the tensorflow weights. This path must not exist yet. - """ - assert tensorflow is not None - tf_major_ver = int(tensorflow.__version__.split(".")[0]) - - if output_path.suffix == ".zip": - output_path = output_path.with_suffix("") - zip_weights = True - else: - zip_weights = False - - if output_path.exists(): - raise ValueError(f"The ouptut directory at {output_path} must not exist.") - - if model.weights.keras_hdf5 is None: - raise ValueError("Missing Keras Hdf5 weights to convert from.") - - weight_spec = model.weights.keras_hdf5 - weight_path = download(weight_spec.source).path - - if weight_spec.tensorflow_version: - model_tf_major_ver = int(weight_spec.tensorflow_version.major) - if model_tf_major_ver != tf_major_ver: - raise RuntimeError( - f"Tensorflow major versions of model {model_tf_major_ver} is not {tf_major_ver}" - ) - - if tf_major_ver == 1: - if len(model.inputs) != 1 or len(model.outputs) != 1: - raise NotImplementedError( - "Weight conversion for models with multiple inputs or outputs is not yet implemented." - ) - return _convert_tf1( - weight_path, - output_path, - model.inputs[0].id, - model.outputs[0].id, - zip_weights, - ) - else: - return _convert_tf2(weight_path, output_path, zip_weights) diff --git a/bioimageio/core/weight_converter/torch/__init__.py b/bioimageio/core/weight_converter/torch/__init__.py deleted file mode 100644 index 1b1ba526..00000000 --- a/bioimageio/core/weight_converter/torch/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# TODO: torch weight converters diff --git a/bioimageio/core/weight_converter/torch/_onnx.py b/bioimageio/core/weight_converter/torch/_onnx.py deleted file mode 100644 index d3c7bf01..00000000 --- a/bioimageio/core/weight_converter/torch/_onnx.py +++ /dev/null @@ -1,100 +0,0 @@ -# type: ignore # TODO: type -from __future__ import annotations -from pathlib import Path -from typing import Any, List, Sequence, cast, Union - -import numpy as np -from numpy.testing import assert_array_almost_equal - -from bioimageio.spec.model import v0_4, v0_5 - -from ...digest_spec import get_member_id, get_test_inputs -from ...weight_converter.torch._utils import load_torch_model - -try: - import torch -except ImportError: - torch = None - - -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 = 15, -) -> v0_5.OnnxWeightsDescr: - """Convert model weights from format 'pytorch_state_dict' to 'onnx'. - - Args: - source_model: model without onnx weights - opset_version: onnx opset version - use_tracing: whether to use tracing or scripting to export the onnx format - test_decimal: precision for testing whether the results agree - """ - state_dict_weights_descr = model_spec.weights.pytorch_state_dict - if state_dict_weights_descr is None: - raise ValueError( - "The provided model does not have weights in the pytorch state dict format" - ) - - assert torch is not None - with torch.no_grad(): - sample = get_test_inputs(model_spec) - 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) - - expected_tensors = model(*input_tensors) - if isinstance(expected_tensors, torch.Tensor): - expected_tensors = [expected_tensors] - expected_outputs: List[np.ndarray[Any, Any]] = [ - out.numpy() for out in expected_tensors - ] - - if use_tracing: - torch.onnx.export( - model, - tuple(input_tensors) if len(input_tensors) > 1 else input_tensors[0], - str(output_path), - verbose=verbose, - opset_version=opset_version, - ) - else: - raise NotImplementedError - - try: - import onnxruntime as rt # pyright: ignore [reportMissingTypeStubs] - except ImportError: - 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)) - onnx_input_node_args = cast( - List[Any], sess.get_inputs() - ) # fixme: remove cast, try using rt.NodeArg instead of Any - onnx_inputs = { - input_name.name: inp - for input_name, inp in zip(onnx_input_node_args, input_data) - } - outputs = cast( - Sequence[np.ndarray[Any, Any]], sess.run(None, onnx_inputs) - ) # FIXME: remove cast - - try: - for exp, out in zip(expected_outputs, outputs): - assert_array_almost_equal(exp, out, decimal=test_decimal) - except AssertionError as e: - 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 - ) diff --git a/bioimageio/core/weight_converter/torch/_torchscript.py b/bioimageio/core/weight_converter/torch/_torchscript.py deleted file mode 100644 index 16dc0128..00000000 --- a/bioimageio/core/weight_converter/torch/_torchscript.py +++ /dev/null @@ -1,145 +0,0 @@ -# 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 -from bioimageio.spec.model.v0_5 import Version - -from ._utils import load_torch_model - -try: - import torch -except ImportError: - torch = None - - -def _check_predictions( - model: Any, - scripted_model: Any, - model_spec: v0_4.ModelDescr | v0_5.ModelDescr, - input_data: Sequence[torch.Tensor], -): - assert torch is not None - - def _check(input_: Sequence[torch.Tensor]) -> None: - expected_tensors = model(*input_) - if isinstance(expected_tensors, torch.Tensor): - expected_tensors = [expected_tensors] - expected_outputs: List[np.ndarray[Any, Any]] = [ - out.numpy() for out in expected_tensors - ] - - output_tensors = scripted_model(*input_) - if isinstance(output_tensors, torch.Tensor): - output_tensors = [output_tensors] - outputs: List[np.ndarray[Any, Any]] = [out.numpy() for out in output_tensors] - - try: - for exp, out in zip(expected_outputs, outputs): - assert_array_almost_equal(exp, out, decimal=4) - except AssertionError as e: - raise ValueError( - f"Results before and after weights conversion do not agree:\n {str(e)}" - ) - - _check(input_data) - - if len(model_spec.inputs) > 1: - return # FIXME: why don't we check multiple inputs? - - input_descr = model_spec.inputs[0] - if isinstance(input_descr, v0_4.InputTensorDescr): - if not isinstance(input_descr.shape, v0_4.ParameterizedInputShape): - return - min_shape = input_descr.shape.min - step = input_descr.shape.step - else: - min_shape: List[int] = [] - step: List[int] = [] - for axis in input_descr.axes: - if isinstance(axis.size, v0_5.ParameterizedSize): - min_shape.append(axis.size.min) - step.append(axis.size.step) - elif isinstance(axis.size, int): - min_shape.append(axis.size) - step.append(0) - elif axis.size is None: - raise NotImplementedError( - f"Can't verify inputs that don't specify their shape fully: {axis}" - ) - elif isinstance(axis.size, v0_5.SizeReference): - raise NotImplementedError(f"Can't handle axes like '{axis}' yet") - else: - assert_never(axis.size) - - 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 step_dim == 0 - else slice(0, max_dim - step_factor * step_dim, 1) - ) - 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( - model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], - output_path: Path, - use_tracing: bool = True, -) -> v0_5.TorchscriptWeightsDescr: - """Convert model weights from format 'pytorch_state_dict' to 'torchscript'. - - Args: - model_descr: location of the resource for the input bioimageio model - 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( - "The provided model does not have weights in the pytorch state dict format" - ) - - input_data = model_descr.get_input_test_arrays() - - 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) - scripted_module: ScriptModule = ( - torch.jit.trace(model, input_data) - if use_tracing - else torch.jit.script(model) - ) - _check_predictions( - model=model, - scripted_model=scripted_module, - model_spec=model_descr, - input_data=input_data, - ) - - scripted_module.save(str(output_path)) - - return v0_5.TorchscriptWeightsDescr( - source=output_path, - pytorch_version=Version(torch.__version__), - parent="pytorch_state_dict", - ) diff --git a/bioimageio/core/weight_converter/torch/_utils.py b/bioimageio/core/weight_converter/torch/_utils.py deleted file mode 100644 index 01df0747..00000000 --- a/bioimageio/core/weight_converter/torch/_utils.py +++ /dev/null @@ -1,24 +0,0 @@ -from typing import Union - -from bioimageio.core.model_adapters._pytorch_model_adapter import PytorchModelAdapter -from bioimageio.spec.model import v0_4, v0_5 -from bioimageio.spec.utils import download - -try: - import torch -except ImportError: - torch = None - - -# additional convenience for pytorch state dict, eventually we want this in python-bioimageio too -# and for each weight format -def load_torch_model( # pyright: ignore[reportUnknownParameterType] - node: Union[v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr], -): - assert torch is not None - model = ( # pyright: ignore[reportUnknownVariableType] - PytorchModelAdapter.get_network(node) - ) - state = torch.load(download(node.source).path, map_location="cpu") - model.load_state_dict(state) # FIXME: check incompatible keys? - return model.eval() # pyright: ignore[reportUnknownVariableType] diff --git a/bioimageio/core/weight_converters.py b/bioimageio/core/weight_converters.py new file mode 100644 index 00000000..6e0d06ec --- /dev/null +++ b/bioimageio/core/weight_converters.py @@ -0,0 +1,492 @@ +# type: ignore # TODO: type +from __future__ import annotations + +import abc +from bioimageio.spec.model.v0_5 import WeightsEntryDescrBase +from typing import Any, List, Sequence, cast, Union +from typing_extensions import assert_never +import numpy as np +from numpy.testing import assert_array_almost_equal +from bioimageio.spec.model import v0_4, v0_5 +from torch.jit import ScriptModule +from bioimageio.core.digest_spec import get_test_inputs, get_member_id +from bioimageio.core.model_adapters._pytorch_model_adapter import PytorchModelAdapter +import os +import shutil +from pathlib import Path +from typing import no_type_check +from zipfile import ZipFile +from bioimageio.spec._internal.version_type import Version +from bioimageio.spec._internal.io_utils import download + +try: + import torch +except ImportError: + torch = None + +try: + import tensorflow.saved_model +except Exception: + tensorflow = None + + +# additional convenience for pytorch state dict, eventually we want this in python-bioimageio too +# and for each weight format +def load_torch_model( # pyright: ignore[reportUnknownParameterType] + node: Union[v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr], +): + assert torch is not None + model = ( # pyright: ignore[reportUnknownVariableType] + PytorchModelAdapter.get_network(node) + ) + state = torch.load(download(node.source).path, map_location="cpu") + model.load_state_dict(state) # FIXME: check incompatible keys? + return model.eval() # pyright: ignore[reportUnknownVariableType] + + +class WeightConverter(abc.ABC): + @abc.abstractmethod + def convert( + self, model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], output_path: Path + ) -> WeightsEntryDescrBase: + raise NotImplementedError + + +class Pytorch2Onnx(WeightConverter): + def __init__(self): + super().__init__() + assert torch is not None + + def convert( + self, + model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], + output_path: Path, + use_tracing: bool = True, + test_decimal: int = 4, + verbose: bool = False, + opset_version: int = 15, + ) -> v0_5.OnnxWeightsDescr: + """ + Convert model weights from the PyTorch state_dict format to the ONNX format. + + Args: + model_descr (Union[v0_4.ModelDescr, v0_5.ModelDescr]): + The model description object that contains the model and its weights. + output_path (Path): + The file path where the ONNX model will be saved. + use_tracing (bool, optional): + Whether to use tracing or scripting to export the ONNX format. Defaults to True. + test_decimal (int, optional): + The decimal precision for comparing the results between the original and converted models. + This is used in the `assert_array_almost_equal` function to check if the outputs match. + Defaults to 4. + verbose (bool, optional): + If True, will print out detailed information during the ONNX export process. Defaults to False. + opset_version (int, optional): + The ONNX opset version to use for the export. Defaults to 15. + + Raises: + ValueError: + If the provided model does not have weights in the PyTorch state_dict format. + ImportError: + If ONNX Runtime is not available for checking the exported ONNX model. + ValueError: + If the results before and after weights conversion do not agree. + + Returns: + v0_5.OnnxWeightsDescr: + A descriptor object that contains information about the exported ONNX weights. + """ + + state_dict_weights_descr = model_descr.weights.pytorch_state_dict + if state_dict_weights_descr is None: + raise ValueError( + "The provided model does not have weights in the pytorch state dict format" + ) + + assert torch is not None + with torch.no_grad(): + sample = get_test_inputs(model_descr) + input_data = [ + sample.members[get_member_id(ipt)].data.data + for ipt in model_descr.inputs + ] + input_tensors = [torch.from_numpy(ipt) for ipt in input_data] + model = load_torch_model(state_dict_weights_descr) + + expected_tensors = model(*input_tensors) + if isinstance(expected_tensors, torch.Tensor): + expected_tensors = [expected_tensors] + expected_outputs: List[np.ndarray[Any, Any]] = [ + out.numpy() for out in expected_tensors + ] + + if use_tracing: + torch.onnx.export( + model, + ( + tuple(input_tensors) + if len(input_tensors) > 1 + else input_tensors[0] + ), + str(output_path), + verbose=verbose, + opset_version=opset_version, + ) + else: + raise NotImplementedError + + try: + import onnxruntime as rt # pyright: ignore [reportMissingTypeStubs] + except ImportError: + 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)) + onnx_input_node_args = cast( + List[Any], sess.get_inputs() + ) # fixme: remove cast, try using rt.NodeArg instead of Any + onnx_inputs = { + input_name.name: inp + for input_name, inp in zip(onnx_input_node_args, input_data) + } + outputs = cast( + Sequence[np.ndarray[Any, Any]], sess.run(None, onnx_inputs) + ) # FIXME: remove cast + + try: + for exp, out in zip(expected_outputs, outputs): + assert_array_almost_equal(exp, out, decimal=test_decimal) + except AssertionError as e: + 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 + ) + + +class Pytorch2Torchscipt(WeightConverter): + def __init__(self): + super().__init__() + assert torch is not None + + def convert( + self, + model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], + output_path: Path, + use_tracing: bool = True, + ) -> v0_5.TorchscriptWeightsDescr: + """ + Convert model weights from the PyTorch `state_dict` format to TorchScript. + + Args: + model_descr (Union[v0_4.ModelDescr, v0_5.ModelDescr]): + The model description object that contains the model and its weights in the PyTorch `state_dict` format. + output_path (Path): + The file path where the TorchScript model will be saved. + use_tracing (bool): + Whether to use tracing or scripting to export the TorchScript format. + - `True`: Use tracing, which is recommended for models with straightforward control flow. + - `False`: Use scripting, which is better for models with dynamic control flow (e.g., loops, conditionals). + + Raises: + ValueError: + If the provided model does not have weights in the PyTorch `state_dict` format. + + Returns: + v0_5.TorchscriptWeightsDescr: + A descriptor object that contains information about the exported TorchScript weights. + """ + state_dict_weights_descr = model_descr.weights.pytorch_state_dict + if state_dict_weights_descr is None: + raise ValueError( + "The provided model does not have weights in the pytorch state dict format" + ) + + input_data = model_descr.get_input_test_arrays() + + 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) + scripted_module: ScriptModule = ( + torch.jit.trace(model, input_data) + if use_tracing + else torch.jit.script(model) + ) + self._check_predictions( + model=model, + scripted_model=scripted_module, + model_spec=model_descr, + input_data=input_data, + ) + + scripted_module.save(str(output_path)) + + return v0_5.TorchscriptWeightsDescr( + source=output_path, + pytorch_version=Version(torch.__version__), + parent="pytorch_state_dict", + ) + + def _check_predictions( + self, + model: Any, + scripted_model: Any, + model_spec: v0_4.ModelDescr | v0_5.ModelDescr, + input_data: Sequence[torch.Tensor], + ): + assert torch is not None + + def _check(input_: Sequence[torch.Tensor]) -> None: + expected_tensors = model(*input_) + if isinstance(expected_tensors, torch.Tensor): + expected_tensors = [expected_tensors] + expected_outputs: List[np.ndarray[Any, Any]] = [ + out.numpy() for out in expected_tensors + ] + + output_tensors = scripted_model(*input_) + if isinstance(output_tensors, torch.Tensor): + output_tensors = [output_tensors] + outputs: List[np.ndarray[Any, Any]] = [ + out.numpy() for out in output_tensors + ] + + try: + for exp, out in zip(expected_outputs, outputs): + assert_array_almost_equal(exp, out, decimal=4) + except AssertionError as e: + raise ValueError( + f"Results before and after weights conversion do not agree:\n {str(e)}" + ) + + _check(input_data) + + if len(model_spec.inputs) > 1: + return # FIXME: why don't we check multiple inputs? + + input_descr = model_spec.inputs[0] + if isinstance(input_descr, v0_4.InputTensorDescr): + if not isinstance(input_descr.shape, v0_4.ParameterizedInputShape): + return + min_shape = input_descr.shape.min + step = input_descr.shape.step + else: + min_shape: List[int] = [] + step: List[int] = [] + for axis in input_descr.axes: + if isinstance(axis.size, v0_5.ParameterizedSize): + min_shape.append(axis.size.min) + step.append(axis.size.step) + elif isinstance(axis.size, int): + min_shape.append(axis.size) + step.append(0) + elif axis.size is None: + raise NotImplementedError( + f"Can't verify inputs that don't specify their shape fully: {axis}" + ) + elif isinstance(axis.size, v0_5.SizeReference): + raise NotImplementedError(f"Can't handle axes like '{axis}' yet") + else: + assert_never(axis.size) + + 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 step_dim == 0 + else slice(0, max_dim - step_factor * step_dim, 1) + ) + 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]) + + +class Tensorflow2Bundled(WeightConverter): + def __init__(self): + super().__init__() + assert tensorflow is not None + + def convert( + self, model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], output_path: Path + ) -> v0_5.TensorflowSavedModelBundleWeightsDescr: + """ + Convert model weights from the 'keras_hdf5' format to the 'tensorflow_saved_model_bundle' format. + + This method handles the conversion of Keras HDF5 model weights into a TensorFlow SavedModel bundle, + which is the recommended format for deploying TensorFlow models. The method supports both TensorFlow 1.x + and 2.x versions, with appropriate checks to ensure compatibility. + + Adapted from: + https://github.com/deepimagej/pydeepimagej/blob/5aaf0e71f9b04df591d5ca596f0af633a7e024f5/pydeepimagej/yaml/create_config.py + + Args: + model_descr (Union[v0_4.ModelDescr, v0_5.ModelDescr]): + The bioimage.io model description containing the model's metadata and weights. + output_path (Path): + The directory where the TensorFlow SavedModel bundle will be saved. + This path must not already exist and, if necessary, will be zipped into a .zip file. + use_tracing (bool): + Placeholder argument; currently not used in this method but required to match the abstract method signature. + + Raises: + ValueError: + - If the specified `output_path` already exists. + - If the Keras HDF5 weights are missing in the model description. + RuntimeError: + If there is a mismatch between the TensorFlow version used by the model and the version installed. + NotImplementedError: + If the model has multiple inputs or outputs and TensorFlow 1.x is being used. + + Returns: + v0_5.TensorflowSavedModelBundleWeightsDescr: + A descriptor object containing information about the converted TensorFlow SavedModel bundle. + """ + assert tensorflow is not None + tf_major_ver = int(tensorflow.__version__.split(".")[0]) + + if output_path.suffix == ".zip": + output_path = output_path.with_suffix("") + zip_weights = True + else: + zip_weights = False + + if output_path.exists(): + raise ValueError(f"The ouptut directory at {output_path} must not exist.") + + if model_descr.weights.keras_hdf5 is None: + raise ValueError("Missing Keras Hdf5 weights to convert from.") + + weight_spec = model_descr.weights.keras_hdf5 + weight_path = download(weight_spec.source).path + + if weight_spec.tensorflow_version: + model_tf_major_ver = int(weight_spec.tensorflow_version.major) + if model_tf_major_ver != tf_major_ver: + raise RuntimeError( + f"Tensorflow major versions of model {model_tf_major_ver} is not {tf_major_ver}" + ) + + if tf_major_ver == 1: + if len(model_descr.inputs) != 1 or len(model_descr.outputs) != 1: + raise NotImplementedError( + "Weight conversion for models with multiple inputs or outputs is not yet implemented." + ) + return self._convert_tf1( + weight_path, + output_path, + model_descr.inputs[0].id, + model_descr.outputs[0].id, + zip_weights, + ) + else: + return self._convert_tf2(weight_path, output_path, zip_weights) + + def _convert_tf2( + self, 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 tensorflow import keras + except Exception: + # if the above fails try to export with the standalone keras + import keras + + model = keras.models.load_model(keras_weight_path) + keras.models.save_model(model, output_path) + + if zip_weights: + output_path = self._zip_model_bundle(output_path) + print("TensorFlow model exported to", output_path) + + return v0_5.TensorflowSavedModelBundleWeightsDescr( + source=output_path, + parent="keras_hdf5", + tensorflow_version=Version(tensorflow.__version__), + ) + + # adapted from + # https://github.com/deepimagej/pydeepimagej/blob/master/pydeepimagej/yaml/create_config.py#L236 + def _convert_tf1( + self, + keras_weight_path: Path, + output_path: Path, + 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 tensorflow import ( + keras, # type: ignore + ) + + except Exception: + # if the above fails try to export with the standalone keras + import keras + + @no_type_check + def build_tf_model(): + keras_model = keras.models.load_model(keras_weight_path) + assert tensorflow is not None + builder = tensorflow.saved_model.builder.SavedModelBuilder(output_path) + signature = ( + tensorflow.saved_model.signature_def_utils.predict_signature_def( + inputs={input_name: keras_model.input}, + outputs={output_name: keras_model.output}, + ) + ) + + signature_def_map = { + tensorflow.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature + } + + builder.add_meta_graph_and_variables( + keras.backend.get_session(), + [tensorflow.saved_model.tag_constants.SERVING], + signature_def_map=signature_def_map, + ) + builder.save() + + build_tf_model() + + if zip_weights: + output_path = self._zip_model_bundle(output_path) + print("TensorFlow model exported to", output_path) + + return v0_5.TensorflowSavedModelBundleWeightsDescr( + source=output_path, + parent="keras_hdf5", + tensorflow_version=Version(tensorflow.__version__), + ) + + def _zip_model_bundle(self, model_bundle_folder: Path): + zipped_model_bundle = model_bundle_folder.with_suffix(".zip") + + with ZipFile(zipped_model_bundle, "w") as zip_obj: + for root, _, files in os.walk(model_bundle_folder): + for filename in files: + src = os.path.join(root, filename) + zip_obj.write(src, os.path.relpath(src, model_bundle_folder)) + + try: + shutil.rmtree(model_bundle_folder) + except Exception: + print("TensorFlow bundled model was not removed after compression") + + return zipped_model_bundle diff --git a/tests/weight_converter/test_add_weights.py b/tests/test_add_weights.py similarity index 100% rename from tests/weight_converter/test_add_weights.py rename to tests/test_add_weights.py diff --git a/tests/test_weight_converters.py b/tests/test_weight_converters.py new file mode 100644 index 00000000..88010744 --- /dev/null +++ b/tests/test_weight_converters.py @@ -0,0 +1,69 @@ +# type: ignore # TODO enable type checking +import zipfile +from pathlib import Path + +import pytest + +import os + +from bioimageio.spec import load_description +from bioimageio.spec.model import v0_5 + +from bioimageio.core.weight_converters import ( + Pytorch2Torchscipt, + Pytorch2Onnx, + Tensorflow2Bundled, +) + + +def test_torchscript_converter(any_torch_model, tmp_path): + bio_model = load_description(any_torch_model) + out_path = tmp_path / "weights.pt" + util = Pytorch2Torchscipt() + ret_val = util.convert(bio_model, out_path) + assert out_path.exists() + assert isinstance(ret_val, v0_5.TorchscriptWeightsDescr) + assert ret_val.source == out_path + + +def test_onnx_converter(convert_to_onnx, tmp_path): + bio_model = load_description(convert_to_onnx) + out_path = tmp_path / "weights.onnx" + opset_version = 15 + util = Pytorch2Onnx() + ret_val = util.convert( + model_descr=bio_model, + output_path=out_path, + test_decimal=3, + opset_version=opset_version, + ) + assert os.path.exists(out_path) + assert isinstance(ret_val, v0_5.OnnxWeightsDescr) + assert ret_val.opset_version == opset_version + assert ret_val.source == out_path + + +def test_tensorflow_converter(any_keras_model: Path, tmp_path: Path): + model = load_description(any_keras_model) + out_path = tmp_path / "weights.h5" + util = Tensorflow2Bundled() + ret_val = util.convert(model, out_path) + assert out_path.exists() + assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr) + assert ret_val.source == out_path + + +@pytest.mark.skip() +def test_tensorflow_converter_zipped(any_keras_model: Path, tmp_path: Path): + out_path = tmp_path / "weights.zip" + model = load_description(any_keras_model) + util = Tensorflow2Bundled() + ret_val = util.convert(model, out_path) + + assert out_path.exists() + assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr) + + expected_names = {"saved_model.pb", "variables/variables.index"} + with zipfile.ZipFile(out_path, "r") as f: + names = set([name for name in f.namelist()]) + assert len(expected_names - names) == 0 diff --git a/tests/weight_converter/keras/test_tensorflow.py b/tests/weight_converter/keras/test_tensorflow.py deleted file mode 100644 index 18a4f2dc..00000000 --- a/tests/weight_converter/keras/test_tensorflow.py +++ /dev/null @@ -1,36 +0,0 @@ -# type: ignore # TODO enable type checking -import zipfile -from pathlib import Path - -import pytest -from bioimageio.spec import load_description -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() -def test_tensorflow_converter(any_keras_model: Path, tmp_path: Path): - model = load_description(any_keras_model) - out_path = tmp_path / "weights.h5" - ret_val = convert_weights_to_tensorflow_saved_model_bundle(model, out_path) - assert out_path.exists() - assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr) - assert ret_val.source == out_path - - -@pytest.mark.skip() -def test_tensorflow_converter_zipped(any_keras_model: Path, tmp_path: Path): - out_path = tmp_path / "weights.zip" - model = load_description(any_keras_model) - ret_val = convert_weights_to_tensorflow_saved_model_bundle(model, out_path) - - assert out_path.exists() - assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr) - - expected_names = {"saved_model.pb", "variables/variables.index"} - with zipfile.ZipFile(out_path, "r") as f: - names = set([name for name in f.namelist()]) - assert len(expected_names - names) == 0 diff --git a/tests/weight_converter/torch/test_onnx.py b/tests/weight_converter/torch/test_onnx.py deleted file mode 100644 index faab39d6..00000000 --- a/tests/weight_converter/torch/test_onnx.py +++ /dev/null @@ -1,23 +0,0 @@ -# type: ignore # TODO enable type checking -import os - -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 - - -def test_onnx_converter(convert_to_onnx, tmp_path): - bio_model = load_description(convert_to_onnx) - out_path = tmp_path / "weights.onnx" - 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) - assert isinstance(ret_val, v0_5.OnnxWeightsDescr) - assert ret_val.opset_version == opset_version - assert ret_val.source == out_path diff --git a/tests/weight_converter/torch/test_torchscript.py b/tests/weight_converter/torch/test_torchscript.py deleted file mode 100644 index 6b397f08..00000000 --- a/tests/weight_converter/torch/test_torchscript.py +++ /dev/null @@ -1,18 +0,0 @@ -# type: ignore # TODO enable type checking -import pytest -from bioimageio.spec import load_description -from bioimageio.spec.model import v0_5 - -from bioimageio.core.weight_converter.torch._torchscript 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(bio_model, out_path) - assert out_path.exists() - assert isinstance(ret_val, v0_5.TorchscriptWeightsDescr) - assert ret_val.source == out_path From db891ebb2456ab9db4696b8753e7bf0e99979d54 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 2 Dec 2024 13:33:09 +0100 Subject: [PATCH 03/47] fix import_callable annotation --- bioimageio/core/digest_spec.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/bioimageio/core/digest_spec.py b/bioimageio/core/digest_spec.py index edb5a45d..854e6a7c 100644 --- a/bioimageio/core/digest_spec.py +++ b/bioimageio/core/digest_spec.py @@ -50,7 +50,12 @@ def import_callable( - node: Union[CallableFromDepencency, ArchitectureFromLibraryDescr], + node: Union[ + ArchitectureFromFileDescr, + ArchitectureFromLibraryDescr, + CallableFromDepencency, + CallableFromFile, + ], /, **kwargs: Unpack[HashKwargs], ) -> Callable[..., Any]: @@ -65,7 +70,6 @@ def import_callable( c = _import_from_file_impl(node.source_file, str(node.callable_name), **kwargs) elif isinstance(node, ArchitectureFromFileDescr): c = _import_from_file_impl(node.source, str(node.callable), sha256=node.sha256) - else: assert_never(node) From a391d940a93566b2350a44aaa2bf61b5fd0d0921 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 2 Dec 2024 13:33:38 +0100 Subject: [PATCH 04/47] improve error traceback for single weights format attempt --- bioimageio/core/model_adapters/_model_adapter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bioimageio/core/model_adapters/_model_adapter.py b/bioimageio/core/model_adapters/_model_adapter.py index c918603e..da2a2ea9 100644 --- a/bioimageio/core/model_adapters/_model_adapter.py +++ b/bioimageio/core/model_adapters/_model_adapter.py @@ -137,7 +137,7 @@ def create( raise ValueError( f"The '{weight_format_priority_order[0]}' model adapter could not be created" + f" in this environment:\n{errors[0][1].__class__.__name__}({errors[0][1]}).\n\n" - ) + ) from errors[0][1] else: error_list = "\n - ".join( From 4103b511e8edf67e5d207613f1555939d7c5e24d Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 2 Dec 2024 15:26:48 +0100 Subject: [PATCH 05/47] add load_state --- .../model_adapters/_pytorch_model_adapter.py | 123 +++++++++++------- .../core/weight_converter/torch/_utils.py | 24 ---- 2 files changed, 73 insertions(+), 74 deletions(-) delete mode 100644 bioimageio/core/weight_converter/torch/_utils.py diff --git a/bioimageio/core/model_adapters/_pytorch_model_adapter.py b/bioimageio/core/model_adapters/_pytorch_model_adapter.py index a5178d74..1992f406 100644 --- a/bioimageio/core/model_adapters/_pytorch_model_adapter.py +++ b/bioimageio/core/model_adapters/_pytorch_model_adapter.py @@ -1,23 +1,23 @@ import gc import warnings -from typing import Any, List, Optional, Sequence, Tuple, Union +from contextlib import nullcontext +from io import TextIOWrapper +from pathlib import Path +from typing import Any, List, Literal, Optional, Sequence, Tuple, Union +import torch +from loguru import logger +from torch import nn +from typing_extensions import assert_never + +from bioimageio.spec.common import ZipPath from bioimageio.spec.model import v0_4, v0_5 from bioimageio.spec.utils import download -from ..axis import AxisId from ..digest_spec import get_axes_infos, import_callable from ..tensor import Tensor from ._model_adapter import ModelAdapter -try: - import torch -except Exception as e: - torch = None - torch_error = str(e) -else: - torch_error = None - class PytorchModelAdapter(ModelAdapter): def __init__( @@ -29,48 +29,41 @@ def __init__( weights: Union[ v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr ], - devices: Optional[Sequence[str]] = None, + devices: Optional[Sequence[Union[str, torch.device]]] = None, + mode: Literal["eval", "train"] = "eval", ): - if torch is None: - raise ImportError(f"failed to import torch: {torch_error}") - super().__init__() self.output_dims = [tuple(a.id for a in get_axes_infos(out)) for out in outputs] - self._network = self.get_network(weights) - self._devices = self.get_devices(devices) - self._network = self._network.to(self._devices[0]) - - self._primary_device = self._devices[0] - state: Any = torch.load( - download(weights).path, - map_location=self._primary_device, # pyright: ignore[reportUnknownArgumentType] - ) - self._network.load_state_dict(state) + devices = self.get_devices(devices) + self._network = self.get_network(weights, load_state=True, devices=devices) + if mode == "eval": + self._network = self._network.eval() + elif mode == "train": + self._network = self._network.train() + else: + assert_never(mode) - self._network = self._network.eval() + self._mode: Literal["eval", "train"] = mode + self._primary_device = devices[0] def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: - if torch is None: - raise ImportError("torch") - with torch.no_grad(): + if self._mode == "eval": + ctxt = torch.no_grad + elif self._mode == "train": + ctxt = nullcontext + else: + assert_never(self._mode) + + with ctxt(): tensors = [ None if ipt is None else torch.from_numpy(ipt.data.data) for ipt in input_tensors ] tensors = [ - ( - None - if t is None - else t.to( - self._primary_device # pyright: ignore[reportUnknownArgumentType] - ) - ) - for t in tensors + (None if t is None else t.to(self._primary_device)) for t in tensors ] result: Union[Tuple[Any, ...], List[Any], Any] - result = self._network( # pyright: ignore[reportUnknownVariableType] - *tensors - ) + result = self._network(*tensors) if not isinstance(result, (tuple, list)): result = [result] @@ -98,14 +91,16 @@ def unload(self) -> None: assert torch is not None torch.cuda.empty_cache() # release reserved memory - @staticmethod - def get_network( # pyright: ignore[reportUnknownParameterType] + @classmethod + def get_network( + cls, weight_spec: Union[ v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr ], - ) -> "torch.nn.Module": # pyright: ignore[reportInvalidTypeForm] - if torch is None: - raise ImportError("torch") + *, + load_state: bool = False, + devices: Optional[Sequence[Union[str, torch.device]]] = None, + ) -> nn.Module: arch = import_callable( weight_spec.architecture, sha256=( @@ -120,19 +115,47 @@ def get_network( # pyright: ignore[reportUnknownParameterType] else weight_spec.architecture.kwargs ) network = arch(**model_kwargs) - if not isinstance(network, torch.nn.Module): + if not isinstance(network, nn.Module): raise ValueError( f"calling {weight_spec.architecture.callable} did not return a torch.nn.Module" ) + if load_state or devices: + use_devices = cls.get_devices(devices) + network = network.to(use_devices[0]) + if load_state: + network = cls.load_state( + network, + path=download(weight_spec).path, + devices=use_devices, + ) + return network + + @staticmethod + def load_state( + network: nn.Module, + path: Union[Path, ZipPath], + devices: Sequence[torch.device], + ) -> nn.Module: + network = network.to(devices[0]) + with path.open("rb") as f: + assert not isinstance(f, TextIOWrapper) + state = torch.load(f, map_location=devices[0]) + + incompatible = network.load_state_dict(state) + if incompatible.missing_keys: + logger.warning("Missing state dict keys: {}", incompatible.missing_keys) + + if incompatible.unexpected_keys: + logger.warning( + "Unexpected state dict keys: {}", incompatible.unexpected_keys + ) return network @staticmethod - def get_devices( # pyright: ignore[reportUnknownParameterType] - devices: Optional[Sequence[str]] = None, - ) -> List["torch.device"]: # pyright: ignore[reportInvalidTypeForm] - if torch is None: - raise ImportError("torch") + def get_devices( + devices: Optional[Sequence[Union[torch.device, str]]] = None, + ) -> List[torch.device]: if not devices: torch_devices = [ ( diff --git a/bioimageio/core/weight_converter/torch/_utils.py b/bioimageio/core/weight_converter/torch/_utils.py deleted file mode 100644 index 01df0747..00000000 --- a/bioimageio/core/weight_converter/torch/_utils.py +++ /dev/null @@ -1,24 +0,0 @@ -from typing import Union - -from bioimageio.core.model_adapters._pytorch_model_adapter import PytorchModelAdapter -from bioimageio.spec.model import v0_4, v0_5 -from bioimageio.spec.utils import download - -try: - import torch -except ImportError: - torch = None - - -# additional convenience for pytorch state dict, eventually we want this in python-bioimageio too -# and for each weight format -def load_torch_model( # pyright: ignore[reportUnknownParameterType] - node: Union[v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr], -): - assert torch is not None - model = ( # pyright: ignore[reportUnknownVariableType] - PytorchModelAdapter.get_network(node) - ) - state = torch.load(download(node.source).path, map_location="cpu") - model.load_state_dict(state) # FIXME: check incompatible keys? - return model.eval() # pyright: ignore[reportUnknownVariableType] From 7ec7afb49a187646b854ba2d94d796d71b85b336 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 2 Dec 2024 15:27:15 +0100 Subject: [PATCH 06/47] update ONNXModelAdapter --- .../model_adapters/_onnx_model_adapter.py | 23 +++++-------------- 1 file changed, 6 insertions(+), 17 deletions(-) diff --git a/bioimageio/core/model_adapters/_onnx_model_adapter.py b/bioimageio/core/model_adapters/_onnx_model_adapter.py index c747de22..87045897 100644 --- a/bioimageio/core/model_adapters/_onnx_model_adapter.py +++ b/bioimageio/core/model_adapters/_onnx_model_adapter.py @@ -1,8 +1,9 @@ import warnings from typing import Any, List, Optional, Sequence, Union -from numpy.typing import NDArray +import onnxruntime as rt +from bioimageio.spec._internal.type_guards import is_list, is_tuple from bioimageio.spec.model import v0_4, v0_5 from bioimageio.spec.utils import download @@ -10,14 +11,6 @@ from ..tensor import Tensor from ._model_adapter import ModelAdapter -try: - import onnxruntime as rt -except Exception as e: - rt = None - rt_error = str(e) -else: - rt_error = None - class ONNXModelAdapter(ModelAdapter): def __init__( @@ -26,9 +19,6 @@ def __init__( model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], devices: Optional[Sequence[str]] = None, ): - if rt is None: - raise ImportError(f"failed to import onnxruntime: {rt_error}") - super().__init__() self._internal_output_axes = [ tuple(a.id for a in get_axes_infos(out)) @@ -51,14 +41,13 @@ def __init__( def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: assert len(input_tensors) == len(self._input_names) input_arrays = [None if ipt is None else ipt.data.data for ipt in input_tensors] - result: Union[Sequence[Optional[NDArray[Any]]], Optional[NDArray[Any]]] - result = self._session.run( # pyright: ignore[reportUnknownVariableType] + result: Any = self._session.run( None, dict(zip(self._input_names, input_arrays)) ) - if isinstance(result, (list, tuple)): - result_seq: Sequence[Optional[NDArray[Any]]] = result + if is_list(result) or is_tuple(result): + result_seq = result else: - result_seq = [result] # type: ignore + result_seq = [result] return [ None if r is None else Tensor(r, dims=axes) From ed8f1db75870af1a2e47ff3842f4a05cdcc217d0 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 2 Dec 2024 15:27:54 +0100 Subject: [PATCH 07/47] update TorchscriptModelAdapter typing --- .../_torchscript_model_adapter.py | 33 +++++-------------- 1 file changed, 8 insertions(+), 25 deletions(-) diff --git a/bioimageio/core/model_adapters/_torchscript_model_adapter.py b/bioimageio/core/model_adapters/_torchscript_model_adapter.py index 0e9f3aef..346718a9 100644 --- a/bioimageio/core/model_adapters/_torchscript_model_adapter.py +++ b/bioimageio/core/model_adapters/_torchscript_model_adapter.py @@ -1,10 +1,10 @@ import gc import warnings -from typing import Any, List, Optional, Sequence, Tuple, Union +from typing import Any, List, Optional, Sequence, Union -import numpy as np -from numpy.typing import NDArray +import torch +from bioimageio.spec._internal.type_guards import is_list, is_ndarray, is_tuple from bioimageio.spec.model import v0_4, v0_5 from bioimageio.spec.utils import download @@ -12,14 +12,6 @@ from ..tensor import Tensor from ._model_adapter import ModelAdapter -try: - import torch -except Exception as e: - torch = None - torch_error = str(e) -else: - torch_error = None - class TorchscriptModelAdapter(ModelAdapter): def __init__( @@ -28,9 +20,6 @@ def __init__( model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], devices: Optional[Sequence[str]] = None, ): - if torch is None: - raise ImportError(f"failed to import torch: {torch_error}") - super().__init__() if model_description.weights.torchscript is None: raise ValueError( @@ -57,19 +46,14 @@ def __init__( ] def forward(self, *batch: Optional[Tensor]) -> List[Optional[Tensor]]: - assert torch is not None with torch.no_grad(): torch_tensor = [ None if b is None else torch.from_numpy(b.data.data).to(self.devices[0]) for b in batch ] - _result: Union[ # pyright: ignore[reportUnknownVariableType] - Tuple[Optional[NDArray[Any]], ...], - List[Optional[NDArray[Any]]], - Optional[NDArray[Any]], - ] = self._model.forward(*torch_tensor) - if isinstance(_result, (tuple, list)): - result: Sequence[Optional[NDArray[Any]]] = _result + _result: Any = self._model.forward(*torch_tensor) + if is_list(_result) or is_tuple(_result): + result: Sequence[Any] = _result else: result = [_result] @@ -77,19 +61,18 @@ def forward(self, *batch: Optional[Tensor]) -> List[Optional[Tensor]]: ( None if r is None - else r.cpu().numpy() if not isinstance(r, np.ndarray) else r + else r.cpu().numpy() if isinstance(r, torch.Tensor) else r ) for r in result ] assert len(result) == len(self._internal_output_axes) return [ - None if r is None else Tensor(r, dims=axes) + None if r is None else Tensor(r, dims=axes) if is_ndarray(r) else r for r, axes in zip(result, self._internal_output_axes) ] def unload(self) -> None: - assert torch is not None self._devices = None del self._model _ = gc.collect() # deallocate memory From 9ae626d3665a977bd55c3cd746c247e08015424b Mon Sep 17 00:00:00 2001 From: fynnbe Date: Tue, 3 Dec 2024 10:47:57 +0100 Subject: [PATCH 08/47] update unzipping in tensorflow model adapter --- .../_tensorflow_model_adapter.py | 40 +++++++++++-------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/bioimageio/core/model_adapters/_tensorflow_model_adapter.py b/bioimageio/core/model_adapters/_tensorflow_model_adapter.py index cfb264f0..b469cde7 100644 --- a/bioimageio/core/model_adapters/_tensorflow_model_adapter.py +++ b/bioimageio/core/model_adapters/_tensorflow_model_adapter.py @@ -1,10 +1,14 @@ import zipfile +from io import TextIOWrapper +from pathlib import Path +from shutil import copyfileobj from typing import List, Literal, Optional, Sequence, Union import numpy as np +import tensorflow as tf # pyright: ignore[reportMissingImports] from loguru import logger -from bioimageio.spec.common import FileSource +from bioimageio.spec.common import FileSource, ZipPath from bioimageio.spec.model import v0_4, v0_5 from bioimageio.spec.utils import download @@ -12,14 +16,6 @@ from ..tensor import Tensor from ._model_adapter import ModelAdapter -try: - import tensorflow as tf # pyright: ignore[reportMissingImports] -except Exception as e: - tf = None - tf_error = str(e) -else: - tf_error = None - class TensorflowModelAdapterBase(ModelAdapter): weight_format: Literal["keras_hdf5", "tensorflow_saved_model_bundle"] @@ -36,9 +32,6 @@ def __init__( ], model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], ): - if tf is None: - raise ImportError(f"failed to import tensorflow: {tf_error}") - super().__init__() self.model_description = model_description tf_version = v0_5.Version( @@ -81,16 +74,29 @@ def __init__( for out in model_description.outputs ] + # TODO: check how to load tf weights without unzipping def require_unzipped(self, weight_file: FileSource): - loacl_weights_file = download(weight_file).path - if zipfile.is_zipfile(loacl_weights_file): - out_path = loacl_weights_file.with_suffix(".unzipped") - with zipfile.ZipFile(loacl_weights_file, "r") as f: + local_weights_file = download(weight_file).path + if isinstance(local_weights_file, ZipPath): + # weights file is in a bioimageio zip package + out_path = ( + Path("bioimageio_unzipped_tf_weights") / local_weights_file.filename + ) + with local_weights_file.open("rb") as src, out_path.open("wb") as dst: + assert not isinstance(src, TextIOWrapper) + copyfileobj(src, dst) + + local_weights_file = out_path + + if zipfile.is_zipfile(local_weights_file): + # weights file itself is a zipfile + out_path = local_weights_file.with_suffix(".unzipped") + with zipfile.ZipFile(local_weights_file, "r") as f: f.extractall(out_path) return out_path else: - return loacl_weights_file + return local_weights_file def _get_network( # pyright: ignore[reportUnknownParameterType] self, weight_file: FileSource From fceed3cd253ba9eff72b5c57a90d4514dc5edd45 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Tue, 3 Dec 2024 13:51:48 +0100 Subject: [PATCH 09/47] add upper bounds to dependencies --- dev/env-py38.yaml | 5 +++-- dev/env-tf.yaml | 3 ++- dev/env-wo-python.yaml | 6 +++--- dev/env.yaml | 9 ++++++--- setup.py | 10 ++++------ 5 files changed, 18 insertions(+), 15 deletions(-) diff --git a/dev/env-py38.yaml b/dev/env-py38.yaml index 22353103..2ff00627 100644 --- a/dev/env-py38.yaml +++ b/dev/env-py38.yaml @@ -3,6 +3,7 @@ name: core38 channels: - conda-forge - nodefaults + - pytorch dependencies: - bioimageio.spec>=0.5.3.5 - black @@ -12,7 +13,7 @@ dependencies: - imageio>=2.5 - jupyter - jupyter-black - # - keras>=3.0 # removed + - keras>=3.0,<4 # removed - loguru - numpy - onnxruntime @@ -28,7 +29,7 @@ dependencies: - pytest-cov - pytest-xdist - python=3.8 # changed - - pytorch>=2.1 + - pytorch>=2.1,<3 - requests - rich - ruff diff --git a/dev/env-tf.yaml b/dev/env-tf.yaml index 0df6fd07..3874009f 100644 --- a/dev/env-tf.yaml +++ b/dev/env-tf.yaml @@ -3,6 +3,7 @@ name: core-tf # changed channels: - conda-forge - nodefaults + # - pytroch # removed dependencies: - bioimageio.spec>=0.5.3.5 - black @@ -28,7 +29,7 @@ dependencies: - pytest-cov - pytest-xdist # - python=3.9 # removed - # - pytorch>=2.1 # removed + # - pytorch>=2.1,<3 # removed - requests - rich # - ruff # removed diff --git a/dev/env-wo-python.yaml b/dev/env-wo-python.yaml index d8cba289..dc76f005 100644 --- a/dev/env-wo-python.yaml +++ b/dev/env-wo-python.yaml @@ -3,7 +3,7 @@ name: core channels: - conda-forge - nodefaults - - pytorch # added + - pytorch dependencies: - bioimageio.spec>=0.5.3.5 - black @@ -13,7 +13,7 @@ dependencies: - imageio>=2.5 - jupyter - jupyter-black - - keras>=3.0 + - keras>=3.0,<4 - loguru - numpy - onnxruntime @@ -29,7 +29,7 @@ dependencies: - pytest-cov - pytest-xdist # - python=3.9 # removed - - pytorch>=2.1 + - pytorch>=2.1,<3 - requests - rich - ruff diff --git a/dev/env.yaml b/dev/env.yaml index 20d60a18..ed16d72e 100644 --- a/dev/env.yaml +++ b/dev/env.yaml @@ -1,6 +1,8 @@ name: core channels: - conda-forge + - nodefaults + - pytorch dependencies: - bioimageio.spec>=0.5.3.5 - black @@ -12,7 +14,7 @@ dependencies: - jupyter-black - ipykernel - matplotlib - - keras>=3.0 + - keras>=3.0,<4 - loguru - numpy - onnxruntime @@ -27,12 +29,13 @@ dependencies: - pytest - pytest-cov - pytest-xdist - - python=3.9 - - pytorch>=2.1 + - python=3.12 + - pytorch>=2.1,<3 - requests - rich - ruff - ruyaml + - tensorflow>=2,<3 - torchvision - tqdm - typing-extensions diff --git a/setup.py b/setup.py index 99747946..79c3b0c9 100644 --- a/setup.py +++ b/setup.py @@ -45,17 +45,17 @@ ], include_package_data=True, extras_require={ - "pytorch": ["torch>=1.6", "torchvision", "keras>=3.0"], - "tensorflow": ["tensorflow", "keras>=2.15"], + "pytorch": (pytorch_deps := ["torch>=1.6,<3", "torchvision", "keras>=3.0,<4"]), + "tensorflow": ["tensorflow", "keras>=2.15,<4"], "onnx": ["onnxruntime"], - "dev": [ + "dev": pytorch_deps + + [ "black", # "crick", # currently requires python<=3.9 "filelock", "jupyter", "jupyter-black", "matplotlib", - "keras>=3.0", "onnxruntime", "packaging>=17.0", "pre-commit", @@ -65,8 +65,6 @@ "pytest-cov", "pytest-xdist", # parallel pytest "pytest", - "torch>=1.6", - "torchvision", ], }, project_urls={ From 77e1e844dea4e4ab7c97d5b243eb7af8f7f97017 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Tue, 3 Dec 2024 14:01:32 +0100 Subject: [PATCH 10/47] update dev envs --- dev/env-py38.yaml | 6 ++++-- dev/env-tf.yaml | 3 ++- dev/env-wo-python.yaml | 4 +++- dev/env.yaml | 5 ++--- 4 files changed, 11 insertions(+), 7 deletions(-) diff --git a/dev/env-py38.yaml b/dev/env-py38.yaml index 2ff00627..7f0d6da9 100644 --- a/dev/env-py38.yaml +++ b/dev/env-py38.yaml @@ -13,8 +13,9 @@ dependencies: - imageio>=2.5 - jupyter - jupyter-black - - keras>=3.0,<4 # removed + - # keras>=3.0,<4 # removed - loguru + - matplotlib - numpy - onnxruntime - packaging>=17.0 @@ -34,9 +35,10 @@ dependencies: - rich - ruff - ruyaml + # - tensorflow>=2,<3 removed - torchvision - tqdm - typing-extensions - xarray - pip: - - -e .. + - -e --no-deps .. diff --git a/dev/env-tf.yaml b/dev/env-tf.yaml index 3874009f..bd12ca02 100644 --- a/dev/env-tf.yaml +++ b/dev/env-tf.yaml @@ -15,6 +15,7 @@ dependencies: - jupyter-black - keras>=2.15 # changed - loguru + - matplotlib - numpy - onnxruntime - packaging>=17.0 @@ -40,4 +41,4 @@ dependencies: - typing-extensions - xarray - pip: - - -e .. + - -e --no-deps .. diff --git a/dev/env-wo-python.yaml b/dev/env-wo-python.yaml index dc76f005..ff0410d9 100644 --- a/dev/env-wo-python.yaml +++ b/dev/env-wo-python.yaml @@ -15,6 +15,7 @@ dependencies: - jupyter-black - keras>=3.0,<4 - loguru + - matplotlib - numpy - onnxruntime - packaging>=17.0 @@ -34,9 +35,10 @@ dependencies: - rich - ruff - ruyaml + - tensorflow>=2,<3 - torchvision - tqdm - typing-extensions - xarray - pip: - - -e .. + - -e --no-deps .. diff --git a/dev/env.yaml b/dev/env.yaml index ed16d72e..c9b62c50 100644 --- a/dev/env.yaml +++ b/dev/env.yaml @@ -12,10 +12,9 @@ dependencies: - imageio>=2.5 - jupyter - jupyter-black - - ipykernel - - matplotlib - keras>=3.0,<4 - loguru + - matplotlib - numpy - onnxruntime - packaging>=17.0 @@ -41,4 +40,4 @@ dependencies: - typing-extensions - xarray - pip: - - -e .. + - -e --no-deps .. From a131369b8e72217ba5ae1f3f3c71b6ea3f09456e Mon Sep 17 00:00:00 2001 From: fynnbe Date: Tue, 3 Dec 2024 14:19:20 +0100 Subject: [PATCH 11/47] WIP setup run expensive tests --- .github/workflows/build.yaml | 15 +++-- bioimageio/core/test_collection.py | 3 + ...t_prediction_pipeline_device_management.py | 2 +- .../core/utils/testing.py => tests/utils.py | 64 +++++++++++-------- 4 files changed, 51 insertions(+), 33 deletions(-) create mode 100644 bioimageio/core/test_collection.py rename bioimageio/core/utils/testing.py => tests/utils.py (80%) diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index c890e8df..634820ad 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -27,6 +27,9 @@ jobs: strategy: matrix: python-version: ['3.8', '3.9', '3.10', '3.11', '3.12'] + include: + - python-version: '3.12' + run-expensive-tests: true steps: - uses: actions/checkout@v4 - name: Install Conda environment with Micromamba @@ -63,6 +66,8 @@ jobs: run: pytest --disable-pytest-warnings env: BIOIMAGEIO_CACHE_PATH: bioimageio_cache + RUN_EXPENSIVE_TESTS: ${{ matrix.run-expensive-tests && 'true' || 'false' }} + test-spec-main: runs-on: ubuntu-latest @@ -71,7 +76,8 @@ jobs: python-version: ['3.8', '3.12'] include: - python-version: '3.12' - is-dev-version: true + report-coverage: true + run-expensive-tests: true steps: - uses: actions/checkout@v4 - name: Install Conda environment with Micromamba @@ -112,17 +118,18 @@ jobs: run: pytest --disable-pytest-warnings env: BIOIMAGEIO_CACHE_PATH: bioimageio_cache - - if: matrix.is-dev-version && github.event_name == 'pull_request' + RUN_EXPENSIVE_TESTS: ${{ matrix.run-expensive-tests && 'true' || 'false' }} + - if: matrix.report-coverage && github.event_name == 'pull_request' uses: orgoro/coverage@v3.2 with: coverageFile: coverage.xml token: ${{ secrets.GITHUB_TOKEN }} - - if: matrix.is-dev-version && github.ref == 'refs/heads/main' + - if: matrix.report-coverage && github.ref == 'refs/heads/main' run: | pip install genbadge[coverage] genbadge coverage --input-file coverage.xml --output-file ./dist/coverage/coverage-badge.svg coverage html -d dist/coverage - - if: matrix.is-dev-version && github.ref == 'refs/heads/main' + - if: matrix.report-coverage && github.ref == 'refs/heads/main' uses: actions/upload-artifact@v4 with: name: coverage diff --git a/bioimageio/core/test_collection.py b/bioimageio/core/test_collection.py new file mode 100644 index 00000000..a0b3025a --- /dev/null +++ b/bioimageio/core/test_collection.py @@ -0,0 +1,3 @@ +from tests.utils import expensive_test + +@expensive_test diff --git a/tests/test_prediction_pipeline_device_management.py b/tests/test_prediction_pipeline_device_management.py index 0e241df1..aa692356 100644 --- a/tests/test_prediction_pipeline_device_management.py +++ b/tests/test_prediction_pipeline_device_management.py @@ -2,9 +2,9 @@ from numpy.testing import assert_array_almost_equal -from bioimageio.core.utils.testing import skip_on from bioimageio.spec.model.v0_4 import ModelDescr as ModelDescr04 from bioimageio.spec.model.v0_5 import ModelDescr, WeightsFormat +from tests.utils import skip_on class TooFewDevicesException(Exception): diff --git a/bioimageio/core/utils/testing.py b/tests/utils.py similarity index 80% rename from bioimageio/core/utils/testing.py rename to tests/utils.py index acd65d95..9cd7445e 100644 --- a/bioimageio/core/utils/testing.py +++ b/tests/utils.py @@ -1,28 +1,36 @@ -# TODO: move to tests/ -from functools import wraps -from typing import Any, Protocol, Type - - -class test_func(Protocol): - def __call__(*args: Any, **kwargs: Any): ... - - -def skip_on(exception: Type[Exception], reason: str): - """adapted from https://stackoverflow.com/a/63522579""" - import pytest - - # Func below is the real decorator and will receive the test function as param - def decorator_func(f: test_func): - @wraps(f) - def wrapper(*args: Any, **kwargs: Any): - try: - # Try to run the test - return f(*args, **kwargs) - except exception: - # If exception of given type happens - # just swallow it and raise pytest.Skip with given reason - pytest.skip(reason) - - return wrapper - - return decorator_func +import os +from functools import wraps +from typing import Any, Protocol, Type + +import pytest + + +class test_func(Protocol): + def __call__(*args: Any, **kwargs: Any): ... + + +def skip_on(exception: Type[Exception], reason: str): + """adapted from https://stackoverflow.com/a/63522579""" + import pytest + + # Func below is the real decorator and will receive the test function as param + def decorator_func(f: test_func): + @wraps(f) + def wrapper(*args: Any, **kwargs: Any): + try: + # Try to run the test + return f(*args, **kwargs) + except exception: + # If exception of given type happens + # just swallow it and raise pytest.Skip with given reason + pytest.skip(reason) + + return wrapper + + return decorator_func + + +expensive_test = pytest.mark.skipif( + (run := os.getenv("RUN_EXPENSIVE_TESTS")) != "true", + reason="Skipping expensive test (enable by RUN_EXPENSIVE_TESTS='true')", +) From 0888c52049ab62a271d9bd0291fbd61edc577382 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Tue, 3 Dec 2024 16:54:03 +0100 Subject: [PATCH 12/47] WIP resource tests --- bioimageio/core/test_bioimageio_collection.py | 60 +++++++++++++++++++ bioimageio/core/test_collection.py | 3 - tests/utils.py | 7 ++- 3 files changed, 66 insertions(+), 4 deletions(-) create mode 100644 bioimageio/core/test_bioimageio_collection.py delete mode 100644 bioimageio/core/test_collection.py diff --git a/bioimageio/core/test_bioimageio_collection.py b/bioimageio/core/test_bioimageio_collection.py new file mode 100644 index 00000000..4de6a26f --- /dev/null +++ b/bioimageio/core/test_bioimageio_collection.py @@ -0,0 +1,60 @@ +from typing import Any, Collection, Dict, Iterable, Mapping, Tuple + +import pytest +import requests +from pydantic import HttpUrl + +from bioimageio.spec import InvalidDescr +from bioimageio.spec.common import Sha256 +from tests.utils import ParameterSet, expensive_test + +BASE_URL = "https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/" + + +def _get_latest_rdf_sources(): + entries: Any = requests.get(BASE_URL + "all_versions.json").json()["entries"] + ret: Dict[str, Tuple[HttpUrl, Sha256]] = {} + for entry in entries: + version = entry["versions"][0] + ret[f"{entry['concept']}/{version['v']}"] = ( + HttpUrl(version["source"]), + Sha256(version["sha256"]), + ) + + return ret + + +ALL_LATEST_RDF_SOURCES: Mapping[str, Tuple[HttpUrl, Sha256]] = _get_latest_rdf_sources() + + +def yield_bioimageio_yaml_urls() -> Iterable[ParameterSet]: + for descr_url, sha in ALL_LATEST_RDF_SOURCES.values(): + key = ( + str(descr_url) + .replace(BASE_URL, "") + .replace("/files/rdf.yaml", "") + .replace("/files/bioimageio.yaml", "") + ) + yield pytest.param(descr_url, sha, key, id=key) + + +KNOWN_INVALID: Collection[str] = set() + + +@expensive_test +@pytest.mark.parametrize("descr_url,sha,key", list(yield_bioimageio_yaml_urls())) +def test_rdf( + descr_url: HttpUrl, + sha: Sha256, + key: str, +): + if key in KNOWN_INVALID: + pytest.skip("known failure") + + from bioimageio.core import load_description_and_test + + descr = load_description_and_test(descr_url, sha256=sha) + assert not isinstance(descr, InvalidDescr) + assert ( + descr.validation_summary.status == "passed" + ), descr.validation_summary.format() diff --git a/bioimageio/core/test_collection.py b/bioimageio/core/test_collection.py deleted file mode 100644 index a0b3025a..00000000 --- a/bioimageio/core/test_collection.py +++ /dev/null @@ -1,3 +0,0 @@ -from tests.utils import expensive_test - -@expensive_test diff --git a/tests/utils.py b/tests/utils.py index 9cd7445e..3a8e695b 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,10 +1,15 @@ import os from functools import wraps -from typing import Any, Protocol, Type +from typing import Any, Protocol, Sequence, Type import pytest +class ParameterSet(Protocol): + def __init__(self, values: Sequence[Any], marks: Any, id: str) -> None: + super().__init__() + + class test_func(Protocol): def __call__(*args: Any, **kwargs: Any): ... From 40dfe2594199b182d7ad81c120ac4db1cc81360d Mon Sep 17 00:00:00 2001 From: fynnbe Date: Tue, 3 Dec 2024 16:54:24 +0100 Subject: [PATCH 13/47] expose sha256 arg --- bioimageio/core/_resource_tests.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/bioimageio/core/_resource_tests.py b/bioimageio/core/_resource_tests.py index 6ace6d5c..8c04b1b8 100644 --- a/bioimageio/core/_resource_tests.py +++ b/bioimageio/core/_resource_tests.py @@ -14,7 +14,7 @@ load_description, ) from bioimageio.spec._internal.common_nodes import ResourceDescrBase -from bioimageio.spec.common import BioimageioYamlContent, PermissiveFileSource +from bioimageio.spec.common import BioimageioYamlContent, PermissiveFileSource, Sha256 from bioimageio.spec.get_conda_env import get_conda_env from bioimageio.spec.model import v0_4, v0_5 from bioimageio.spec.model.v0_5 import WeightsFormat @@ -102,6 +102,7 @@ def test_model( decimal: Optional[int] = None, *, determinism: Literal["seed_only", "full"] = "seed_only", + sha256: Optional[Sha256] = None, ) -> ValidationSummary: """Test model inference""" return test_description( @@ -113,6 +114,7 @@ def test_model( decimal=decimal, determinism=determinism, expected_type="model", + sha256=sha256, ) @@ -127,6 +129,7 @@ def test_description( decimal: Optional[int] = None, determinism: Literal["seed_only", "full"] = "seed_only", expected_type: Optional[str] = None, + sha256: Optional[Sha256] = None, ) -> ValidationSummary: """Test a bioimage.io resource dynamically, e.g. prediction of test tensors for models""" rd = load_description_and_test( @@ -139,6 +142,7 @@ def test_description( decimal=decimal, determinism=determinism, expected_type=expected_type, + sha256=sha256, ) return rd.validation_summary @@ -154,6 +158,7 @@ def load_description_and_test( decimal: Optional[int] = None, determinism: Literal["seed_only", "full"] = "seed_only", expected_type: Optional[str] = None, + sha256: Optional[Sha256] = None, ) -> Union[ResourceDescr, InvalidDescr]: """Test RDF dynamically, e.g. model inference of test inputs""" if ( @@ -171,7 +176,7 @@ def load_description_and_test( elif isinstance(source, dict): rd = build_description(source, format_version=format_version) else: - rd = load_description(source, format_version=format_version) + rd = load_description(source, format_version=format_version, sha256=sha256) rd.validation_summary.env.add( InstalledPackage(name="bioimageio.core", version=VERSION) From bb539d4d676b19d3ce5371b2d6f8830001424910 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 4 Dec 2024 14:40:09 +0100 Subject: [PATCH 14/47] update torchscript adapter --- .../core/weight_converter/torch/_torchscript.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/bioimageio/core/weight_converter/torch/_torchscript.py b/bioimageio/core/weight_converter/torch/_torchscript.py index 5ca16069..766bc7e3 100644 --- a/bioimageio/core/weight_converter/torch/_torchscript.py +++ b/bioimageio/core/weight_converter/torch/_torchscript.py @@ -1,20 +1,15 @@ -# type: ignore # TODO: type from pathlib import Path from typing import List, Sequence, Union import numpy as np +import torch from numpy.testing import assert_array_almost_equal from typing_extensions import Any, assert_never from bioimageio.spec.model import v0_4, v0_5 from bioimageio.spec.model.v0_5 import Version -from ._utils import load_torch_model - -try: - import torch -except ImportError: - torch = None +from ...model_adapters._pytorch_model_adapter import PytorchModelAdapter # FIXME: remove Any @@ -119,7 +114,9 @@ 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) + model = PytorchModelAdapter.get_network( + state_dict_weights_descr, load_state=True + ) # FIXME: remove Any if use_tracing: From fedd43ce1fd158b1da821b91f4b8292088a226a7 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 4 Dec 2024 14:40:31 +0100 Subject: [PATCH 15/47] bump spec lib version --- bioimageio/core/test_bioimageio_collection.py | 2 +- dev/env-py38.yaml | 2 +- dev/env-tf.yaml | 2 +- dev/env-wo-python.yaml | 2 +- dev/env.yaml | 4 ++-- setup.py | 2 +- 6 files changed, 7 insertions(+), 7 deletions(-) diff --git a/bioimageio/core/test_bioimageio_collection.py b/bioimageio/core/test_bioimageio_collection.py index 4de6a26f..2cf9ced0 100644 --- a/bioimageio/core/test_bioimageio_collection.py +++ b/bioimageio/core/test_bioimageio_collection.py @@ -17,7 +17,7 @@ def _get_latest_rdf_sources(): for entry in entries: version = entry["versions"][0] ret[f"{entry['concept']}/{version['v']}"] = ( - HttpUrl(version["source"]), + HttpUrl(version["source"]), # pyright: ignore[reportCallIssue] Sha256(version["sha256"]), ) diff --git a/dev/env-py38.yaml b/dev/env-py38.yaml index 7f0d6da9..69030cc9 100644 --- a/dev/env-py38.yaml +++ b/dev/env-py38.yaml @@ -5,7 +5,7 @@ channels: - nodefaults - pytorch dependencies: - - bioimageio.spec>=0.5.3.5 + - bioimageio.spec>=0.5.3.6 - black - crick # uncommented - filelock diff --git a/dev/env-tf.yaml b/dev/env-tf.yaml index bd12ca02..799d2a59 100644 --- a/dev/env-tf.yaml +++ b/dev/env-tf.yaml @@ -5,7 +5,7 @@ channels: - nodefaults # - pytroch # removed dependencies: - - bioimageio.spec>=0.5.3.5 + - bioimageio.spec>=0.5.3.6 - black # - crick # currently requires python<=3.9 - filelock diff --git a/dev/env-wo-python.yaml b/dev/env-wo-python.yaml index ff0410d9..a0b7c978 100644 --- a/dev/env-wo-python.yaml +++ b/dev/env-wo-python.yaml @@ -5,7 +5,7 @@ channels: - nodefaults - pytorch dependencies: - - bioimageio.spec>=0.5.3.5 + - bioimageio.spec>=0.5.3.6 - black # - crick # currently requires python<=3.9 - filelock diff --git a/dev/env.yaml b/dev/env.yaml index c9b62c50..a65158d9 100644 --- a/dev/env.yaml +++ b/dev/env.yaml @@ -1,10 +1,10 @@ -name: core +name: full channels: - conda-forge - nodefaults - pytorch dependencies: - - bioimageio.spec>=0.5.3.5 + - bioimageio.spec>=0.5.3.6 - black # - crick # currently requires python<=3.9 - filelock diff --git a/setup.py b/setup.py index 79c3b0c9..af913c1d 100644 --- a/setup.py +++ b/setup.py @@ -30,7 +30,7 @@ ], packages=find_namespace_packages(exclude=["tests"]), install_requires=[ - "bioimageio.spec ==0.5.3.5", + "bioimageio.spec ==0.5.3.6", "h5py", "imageio>=2.10", "loguru", From 7f6fdf1510f66e17e9c698fdb493adcbca8aa173 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 5 Dec 2024 09:27:16 +0100 Subject: [PATCH 16/47] WIP refactor backend libs --- bioimageio/core/_create_model_adapter.py | 127 ++++++++ bioimageio/core/_model_adapter.py | 93 ++++++ bioimageio/core/backend/__init__.py | 0 .../keras.py} | 2 +- bioimageio/core/backend/pytorch.py | 176 +++++++++++ bioimageio/core/model_adapters.py | 8 + bioimageio/core/model_adapters/__init__.py | 7 - .../core/model_adapters/_model_adapter.py | 2 +- .../model_adapters/_onnx_model_adapter.py | 60 ---- .../model_adapters/_pytorch_model_adapter.py | 176 ----------- .../_tensorflow_model_adapter.py | 281 ------------------ .../_torchscript_model_adapter.py | 79 ----- bioimageio/core/weight_converters/__init__.py | 0 .../core/weight_converters/_add_weights.py | 25 ++ .../core/weight_converters/pytorch_to_onnx.py | 124 ++++++++ 15 files changed, 555 insertions(+), 605 deletions(-) create mode 100644 bioimageio/core/_create_model_adapter.py create mode 100644 bioimageio/core/_model_adapter.py create mode 100644 bioimageio/core/backend/__init__.py rename bioimageio/core/{model_adapters/_keras_model_adapter.py => backend/keras.py} (98%) create mode 100644 bioimageio/core/backend/pytorch.py create mode 100644 bioimageio/core/model_adapters.py delete mode 100644 bioimageio/core/model_adapters/__init__.py delete mode 100644 bioimageio/core/model_adapters/_onnx_model_adapter.py delete mode 100644 bioimageio/core/model_adapters/_tensorflow_model_adapter.py delete mode 100644 bioimageio/core/model_adapters/_torchscript_model_adapter.py create mode 100644 bioimageio/core/weight_converters/__init__.py create mode 100644 bioimageio/core/weight_converters/_add_weights.py create mode 100644 bioimageio/core/weight_converters/pytorch_to_onnx.py diff --git a/bioimageio/core/_create_model_adapter.py b/bioimageio/core/_create_model_adapter.py new file mode 100644 index 00000000..ee79f260 --- /dev/null +++ b/bioimageio/core/_create_model_adapter.py @@ -0,0 +1,127 @@ +import warnings +from abc import abstractmethod +from typing import List, Optional, Sequence, Tuple, Union, final + +from bioimageio.spec.model import v0_4, v0_5 + +from ._model_adapter import ( + DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER, + ModelAdapter, + WeightsFormat, +) +from .tensor import Tensor + + +def create_model_adapter( + model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], + *, + devices: Optional[Sequence[str]] = None, + weight_format_priority_order: Optional[Sequence[WeightsFormat]] = None, +): + """ + 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 + """ + if not isinstance(model_description, (v0_4.ModelDescr, v0_5.ModelDescr)): + raise TypeError( + f"expected v0_4.ModelDescr or v0_5.ModelDescr, but got {type(model_description)}" + ) + + weights = model_description.weights + errors: List[Tuple[WeightsFormat, Exception]] = [] + weight_format_priority_order = ( + DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER + if weight_format_priority_order is None + else weight_format_priority_order + ) + # limit weight formats to the ones present + weight_format_priority_order = [ + w for w in weight_format_priority_order if getattr(weights, w) is not None + ] + + for wf in weight_format_priority_order: + if wf == "pytorch_state_dict" and weights.pytorch_state_dict is not None: + try: + from .model_adapters_old._pytorch_model_adapter import ( + PytorchModelAdapter, + ) + + return PytorchModelAdapter( + outputs=model_description.outputs, + weights=weights.pytorch_state_dict, + devices=devices, + ) + except Exception as e: + errors.append((wf, e)) + elif ( + wf == "tensorflow_saved_model_bundle" + and weights.tensorflow_saved_model_bundle is not None + ): + try: + from .model_adapters_old._tensorflow_model_adapter import ( + TensorflowModelAdapter, + ) + + return TensorflowModelAdapter( + model_description=model_description, devices=devices + ) + except Exception as e: + errors.append((wf, e)) + elif wf == "onnx" and weights.onnx is not None: + try: + from .model_adapters_old._onnx_model_adapter import ONNXModelAdapter + + return ONNXModelAdapter( + model_description=model_description, devices=devices + ) + except Exception as e: + errors.append((wf, e)) + elif wf == "torchscript" and weights.torchscript is not None: + try: + from .model_adapters_old._torchscript_model_adapter import ( + TorchscriptModelAdapter, + ) + + return TorchscriptModelAdapter( + model_description=model_description, devices=devices + ) + except Exception as e: + errors.append((wf, e)) + elif wf == "keras_hdf5" and weights.keras_hdf5 is not None: + # keras can either be installed as a separate package or used as part of tensorflow + # we try to first import the keras model adapter using the separate package and, + # if it is not available, try to load the one using tf + try: + from .backend.keras import ( + KerasModelAdapter, + keras, # type: ignore + ) + + if keras is None: + from .model_adapters_old._tensorflow_model_adapter import ( + KerasModelAdapter, + ) + + return KerasModelAdapter( + model_description=model_description, devices=devices + ) + except Exception as e: + errors.append((wf, e)) + + assert errors + if len(weight_format_priority_order) == 1: + assert len(errors) == 1 + raise ValueError( + f"The '{weight_format_priority_order[0]}' model adapter could not be created" + + f" in this environment:\n{errors[0][1].__class__.__name__}({errors[0][1]}).\n\n" + ) from errors[0][1] + + else: + error_list = "\n - ".join( + f"{wf}: {e.__class__.__name__}({e})" for wf, e in errors + ) + raise ValueError( + "None of the weight format specific model adapters could be created" + + f" in this environment. Errors are:\n\n{error_list}.\n\n" + ) diff --git a/bioimageio/core/_model_adapter.py b/bioimageio/core/_model_adapter.py new file mode 100644 index 00000000..0438d35e --- /dev/null +++ b/bioimageio/core/_model_adapter.py @@ -0,0 +1,93 @@ +import warnings +from abc import ABC, abstractmethod +from typing import List, Optional, Sequence, Tuple, Union, final + +from bioimageio.spec.model import v0_4, v0_5 + +from .tensor import Tensor + +WeightsFormat = Union[v0_4.WeightsFormat, v0_5.WeightsFormat] + +__all__ = [ + "ModelAdapter", + "create_model_adapter", + "get_weight_formats", +] + +# Known weight formats in order of priority +# First match wins +DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER: Tuple[WeightsFormat, ...] = ( + "pytorch_state_dict", + "tensorflow_saved_model_bundle", + "torchscript", + "onnx", + "keras_hdf5", +) + + +class ModelAdapter(ABC): + """ + Represents model *without* any preprocessing or postprocessing. + + ``` + from bioimageio.core import load_description + + model = load_description(...) + + # option 1: + adapter = ModelAdapter.create(model) + adapter.forward(...) + adapter.unload() + + # option 2: + with ModelAdapter.create(model) as adapter: + adapter.forward(...) + ``` + """ + + @final + @classmethod + def create( + cls, + model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], + *, + devices: Optional[Sequence[str]] = None, + weight_format_priority_order: Optional[Sequence[WeightsFormat]] = None, + ): + """ + 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 + """ + from ._create_model_adapter import create_model_adapter + + return create_model_adapter( + model_description, + devices=devices, + weight_format_priority_order=weight_format_priority_order, + ) + + @final + def load(self, *, devices: Optional[Sequence[str]] = None) -> None: + warnings.warn("Deprecated. ModelAdapter is loaded on initialization") + + @abstractmethod + def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: + """ + Run forward pass of model to get model predictions + """ + # TODO: handle tensor.transpose in here and make _forward_impl the abstract impl + + @abstractmethod + def unload(self): + """ + Unload model from any devices, freeing their memory. + The moder adapter should be considered unusable afterwards. + """ + + +def get_weight_formats() -> List[str]: + """ + Return list of supported weight types + """ + return list(DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER) diff --git a/bioimageio/core/backend/__init__.py b/bioimageio/core/backend/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bioimageio/core/model_adapters/_keras_model_adapter.py b/bioimageio/core/backend/keras.py similarity index 98% rename from bioimageio/core/model_adapters/_keras_model_adapter.py rename to bioimageio/core/backend/keras.py index e6864ccc..1d273cfc 100644 --- a/bioimageio/core/model_adapters/_keras_model_adapter.py +++ b/bioimageio/core/backend/keras.py @@ -10,8 +10,8 @@ from .._settings import settings from ..digest_spec import get_axes_infos +from ..model_adapters import ModelAdapter from ..tensor import Tensor -from ._model_adapter import ModelAdapter os.environ["KERAS_BACKEND"] = settings.keras_backend diff --git a/bioimageio/core/backend/pytorch.py b/bioimageio/core/backend/pytorch.py new file mode 100644 index 00000000..1992f406 --- /dev/null +++ b/bioimageio/core/backend/pytorch.py @@ -0,0 +1,176 @@ +import gc +import warnings +from contextlib import nullcontext +from io import TextIOWrapper +from pathlib import Path +from typing import Any, List, Literal, Optional, Sequence, Tuple, Union + +import torch +from loguru import logger +from torch import nn +from typing_extensions import assert_never + +from bioimageio.spec.common import ZipPath +from bioimageio.spec.model import v0_4, v0_5 +from bioimageio.spec.utils import download + +from ..digest_spec import get_axes_infos, import_callable +from ..tensor import Tensor +from ._model_adapter import ModelAdapter + + +class PytorchModelAdapter(ModelAdapter): + def __init__( + self, + *, + outputs: Union[ + Sequence[v0_4.OutputTensorDescr], Sequence[v0_5.OutputTensorDescr] + ], + weights: Union[ + v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr + ], + devices: Optional[Sequence[Union[str, torch.device]]] = None, + mode: Literal["eval", "train"] = "eval", + ): + super().__init__() + self.output_dims = [tuple(a.id for a in get_axes_infos(out)) for out in outputs] + devices = self.get_devices(devices) + self._network = self.get_network(weights, load_state=True, devices=devices) + if mode == "eval": + self._network = self._network.eval() + elif mode == "train": + self._network = self._network.train() + else: + assert_never(mode) + + self._mode: Literal["eval", "train"] = mode + self._primary_device = devices[0] + + def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: + if self._mode == "eval": + ctxt = torch.no_grad + elif self._mode == "train": + ctxt = nullcontext + else: + assert_never(self._mode) + + with ctxt(): + tensors = [ + None if ipt is None else torch.from_numpy(ipt.data.data) + for ipt in input_tensors + ] + tensors = [ + (None if t is None else t.to(self._primary_device)) for t in tensors + ] + result: Union[Tuple[Any, ...], List[Any], Any] + result = self._network(*tensors) + if not isinstance(result, (tuple, list)): + result = [result] + + result = [ + ( + None + if r is None + else r.detach().cpu().numpy() if isinstance(r, torch.Tensor) else r + ) + for r in result # pyright: ignore[reportUnknownVariableType] + ] + if len(result) > len(self.output_dims): + raise ValueError( + f"Expected at most {len(self.output_dims)} outputs, but got {len(result)}" + ) + + return [ + None if r is None else Tensor(r, dims=out) + for r, out in zip(result, self.output_dims) + ] + + def unload(self) -> None: + del self._network + _ = gc.collect() # deallocate memory + assert torch is not None + torch.cuda.empty_cache() # release reserved memory + + @classmethod + def get_network( + cls, + weight_spec: Union[ + v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr + ], + *, + load_state: bool = False, + devices: Optional[Sequence[Union[str, torch.device]]] = None, + ) -> nn.Module: + arch = import_callable( + weight_spec.architecture, + sha256=( + weight_spec.architecture_sha256 + if isinstance(weight_spec, v0_4.PytorchStateDictWeightsDescr) + else weight_spec.sha256 + ), + ) + model_kwargs = ( + weight_spec.kwargs + if isinstance(weight_spec, v0_4.PytorchStateDictWeightsDescr) + else weight_spec.architecture.kwargs + ) + network = arch(**model_kwargs) + if not isinstance(network, nn.Module): + raise ValueError( + f"calling {weight_spec.architecture.callable} did not return a torch.nn.Module" + ) + + if load_state or devices: + use_devices = cls.get_devices(devices) + network = network.to(use_devices[0]) + if load_state: + network = cls.load_state( + network, + path=download(weight_spec).path, + devices=use_devices, + ) + return network + + @staticmethod + def load_state( + network: nn.Module, + path: Union[Path, ZipPath], + devices: Sequence[torch.device], + ) -> nn.Module: + network = network.to(devices[0]) + with path.open("rb") as f: + assert not isinstance(f, TextIOWrapper) + state = torch.load(f, map_location=devices[0]) + + incompatible = network.load_state_dict(state) + if incompatible.missing_keys: + logger.warning("Missing state dict keys: {}", incompatible.missing_keys) + + if incompatible.unexpected_keys: + logger.warning( + "Unexpected state dict keys: {}", incompatible.unexpected_keys + ) + return network + + @staticmethod + def get_devices( + devices: Optional[Sequence[Union[torch.device, str]]] = None, + ) -> List[torch.device]: + if not devices: + torch_devices = [ + ( + torch.device("cuda") + if torch.cuda.is_available() + else torch.device("cpu") + ) + ] + else: + torch_devices = [torch.device(d) for d in devices] + + if len(torch_devices) > 1: + warnings.warn( + f"Multiple devices for single pytorch model not yet implemented; ignoring {torch_devices[1:]}" + ) + torch_devices = torch_devices[:1] + + return torch_devices diff --git a/bioimageio/core/model_adapters.py b/bioimageio/core/model_adapters.py new file mode 100644 index 00000000..86fcfe4b --- /dev/null +++ b/bioimageio/core/model_adapters.py @@ -0,0 +1,8 @@ +from ._create_model_adapter import create_model_adapter +from ._model_adapter import ModelAdapter, get_weight_formats + +__all__ = [ + "ModelAdapter", + "create_model_adapter", + "get_weight_formats", +] diff --git a/bioimageio/core/model_adapters/__init__.py b/bioimageio/core/model_adapters/__init__.py deleted file mode 100644 index 01899de9..00000000 --- a/bioimageio/core/model_adapters/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from ._model_adapter import ModelAdapter, create_model_adapter, get_weight_formats - -__all__ = [ - "ModelAdapter", - "create_model_adapter", - "get_weight_formats", -] diff --git a/bioimageio/core/model_adapters/_model_adapter.py b/bioimageio/core/model_adapters/_model_adapter.py index da2a2ea9..3921f81b 100644 --- a/bioimageio/core/model_adapters/_model_adapter.py +++ b/bioimageio/core/model_adapters/_model_adapter.py @@ -117,7 +117,7 @@ def create( # we try to first import the keras model adapter using the separate package and, # if it is not available, try to load the one using tf try: - from ._keras_model_adapter import ( + from ._keras import ( KerasModelAdapter, keras, # type: ignore ) diff --git a/bioimageio/core/model_adapters/_onnx_model_adapter.py b/bioimageio/core/model_adapters/_onnx_model_adapter.py deleted file mode 100644 index 87045897..00000000 --- a/bioimageio/core/model_adapters/_onnx_model_adapter.py +++ /dev/null @@ -1,60 +0,0 @@ -import warnings -from typing import Any, List, Optional, Sequence, Union - -import onnxruntime as rt - -from bioimageio.spec._internal.type_guards import is_list, is_tuple -from bioimageio.spec.model import v0_4, v0_5 -from bioimageio.spec.utils import download - -from ..digest_spec import get_axes_infos -from ..tensor import Tensor -from ._model_adapter import ModelAdapter - - -class ONNXModelAdapter(ModelAdapter): - def __init__( - self, - *, - model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], - devices: Optional[Sequence[str]] = None, - ): - super().__init__() - self._internal_output_axes = [ - tuple(a.id for a in get_axes_infos(out)) - for out in model_description.outputs - ] - if model_description.weights.onnx is None: - raise ValueError("No ONNX weights specified for {model_description.name}") - - self._session = rt.InferenceSession( - str(download(model_description.weights.onnx.source).path) - ) - onnx_inputs = self._session.get_inputs() # type: ignore - self._input_names: List[str] = [ipt.name for ipt in onnx_inputs] # type: ignore - - 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: Optional[Tensor]) -> List[Optional[Tensor]]: - assert len(input_tensors) == len(self._input_names) - input_arrays = [None if ipt is None else ipt.data.data for ipt in input_tensors] - result: Any = self._session.run( - None, dict(zip(self._input_names, input_arrays)) - ) - if is_list(result) or is_tuple(result): - result_seq = result - else: - result_seq = [result] - - return [ - None if r is None else Tensor(r, dims=axes) - for r, axes in zip(result_seq, self._internal_output_axes) - ] - - def unload(self) -> None: - warnings.warn( - "Device management is not implemented for onnx yet, cannot unload model" - ) diff --git a/bioimageio/core/model_adapters/_pytorch_model_adapter.py b/bioimageio/core/model_adapters/_pytorch_model_adapter.py index 1992f406..e69de29b 100644 --- a/bioimageio/core/model_adapters/_pytorch_model_adapter.py +++ b/bioimageio/core/model_adapters/_pytorch_model_adapter.py @@ -1,176 +0,0 @@ -import gc -import warnings -from contextlib import nullcontext -from io import TextIOWrapper -from pathlib import Path -from typing import Any, List, Literal, Optional, Sequence, Tuple, Union - -import torch -from loguru import logger -from torch import nn -from typing_extensions import assert_never - -from bioimageio.spec.common import ZipPath -from bioimageio.spec.model import v0_4, v0_5 -from bioimageio.spec.utils import download - -from ..digest_spec import get_axes_infos, import_callable -from ..tensor import Tensor -from ._model_adapter import ModelAdapter - - -class PytorchModelAdapter(ModelAdapter): - def __init__( - self, - *, - outputs: Union[ - Sequence[v0_4.OutputTensorDescr], Sequence[v0_5.OutputTensorDescr] - ], - weights: Union[ - v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr - ], - devices: Optional[Sequence[Union[str, torch.device]]] = None, - mode: Literal["eval", "train"] = "eval", - ): - super().__init__() - self.output_dims = [tuple(a.id for a in get_axes_infos(out)) for out in outputs] - devices = self.get_devices(devices) - self._network = self.get_network(weights, load_state=True, devices=devices) - if mode == "eval": - self._network = self._network.eval() - elif mode == "train": - self._network = self._network.train() - else: - assert_never(mode) - - self._mode: Literal["eval", "train"] = mode - self._primary_device = devices[0] - - def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: - if self._mode == "eval": - ctxt = torch.no_grad - elif self._mode == "train": - ctxt = nullcontext - else: - assert_never(self._mode) - - with ctxt(): - tensors = [ - None if ipt is None else torch.from_numpy(ipt.data.data) - for ipt in input_tensors - ] - tensors = [ - (None if t is None else t.to(self._primary_device)) for t in tensors - ] - result: Union[Tuple[Any, ...], List[Any], Any] - result = self._network(*tensors) - if not isinstance(result, (tuple, list)): - result = [result] - - result = [ - ( - None - if r is None - else r.detach().cpu().numpy() if isinstance(r, torch.Tensor) else r - ) - for r in result # pyright: ignore[reportUnknownVariableType] - ] - if len(result) > len(self.output_dims): - raise ValueError( - f"Expected at most {len(self.output_dims)} outputs, but got {len(result)}" - ) - - return [ - None if r is None else Tensor(r, dims=out) - for r, out in zip(result, self.output_dims) - ] - - def unload(self) -> None: - del self._network - _ = gc.collect() # deallocate memory - assert torch is not None - torch.cuda.empty_cache() # release reserved memory - - @classmethod - def get_network( - cls, - weight_spec: Union[ - v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr - ], - *, - load_state: bool = False, - devices: Optional[Sequence[Union[str, torch.device]]] = None, - ) -> nn.Module: - arch = import_callable( - weight_spec.architecture, - sha256=( - weight_spec.architecture_sha256 - if isinstance(weight_spec, v0_4.PytorchStateDictWeightsDescr) - else weight_spec.sha256 - ), - ) - model_kwargs = ( - weight_spec.kwargs - if isinstance(weight_spec, v0_4.PytorchStateDictWeightsDescr) - else weight_spec.architecture.kwargs - ) - network = arch(**model_kwargs) - if not isinstance(network, nn.Module): - raise ValueError( - f"calling {weight_spec.architecture.callable} did not return a torch.nn.Module" - ) - - if load_state or devices: - use_devices = cls.get_devices(devices) - network = network.to(use_devices[0]) - if load_state: - network = cls.load_state( - network, - path=download(weight_spec).path, - devices=use_devices, - ) - return network - - @staticmethod - def load_state( - network: nn.Module, - path: Union[Path, ZipPath], - devices: Sequence[torch.device], - ) -> nn.Module: - network = network.to(devices[0]) - with path.open("rb") as f: - assert not isinstance(f, TextIOWrapper) - state = torch.load(f, map_location=devices[0]) - - incompatible = network.load_state_dict(state) - if incompatible.missing_keys: - logger.warning("Missing state dict keys: {}", incompatible.missing_keys) - - if incompatible.unexpected_keys: - logger.warning( - "Unexpected state dict keys: {}", incompatible.unexpected_keys - ) - return network - - @staticmethod - def get_devices( - devices: Optional[Sequence[Union[torch.device, str]]] = None, - ) -> List[torch.device]: - if not devices: - torch_devices = [ - ( - torch.device("cuda") - if torch.cuda.is_available() - else torch.device("cpu") - ) - ] - else: - torch_devices = [torch.device(d) for d in devices] - - if len(torch_devices) > 1: - warnings.warn( - f"Multiple devices for single pytorch model not yet implemented; ignoring {torch_devices[1:]}" - ) - torch_devices = torch_devices[:1] - - return torch_devices diff --git a/bioimageio/core/model_adapters/_tensorflow_model_adapter.py b/bioimageio/core/model_adapters/_tensorflow_model_adapter.py deleted file mode 100644 index b469cde7..00000000 --- a/bioimageio/core/model_adapters/_tensorflow_model_adapter.py +++ /dev/null @@ -1,281 +0,0 @@ -import zipfile -from io import TextIOWrapper -from pathlib import Path -from shutil import copyfileobj -from typing import List, Literal, Optional, Sequence, Union - -import numpy as np -import tensorflow as tf # pyright: ignore[reportMissingImports] -from loguru import logger - -from bioimageio.spec.common import FileSource, ZipPath -from bioimageio.spec.model import v0_4, v0_5 -from bioimageio.spec.utils import download - -from ..digest_spec import get_axes_infos -from ..tensor import Tensor -from ._model_adapter import ModelAdapter - - -class TensorflowModelAdapterBase(ModelAdapter): - weight_format: Literal["keras_hdf5", "tensorflow_saved_model_bundle"] - - def __init__( - self, - *, - devices: Optional[Sequence[str]] = None, - weights: Union[ - v0_4.KerasHdf5WeightsDescr, - v0_4.TensorflowSavedModelBundleWeightsDescr, - v0_5.KerasHdf5WeightsDescr, - v0_5.TensorflowSavedModelBundleWeightsDescr, - ], - model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], - ): - super().__init__() - self.model_description = model_description - tf_version = v0_5.Version( - tf.__version__ # pyright: ignore[reportUnknownArgumentType] - ) - model_tf_version = weights.tensorflow_version - if model_tf_version is None: - 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: - 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, - ): - logger.warning( - "The tensorflow version specified by the model does not match the installed: " - + f"{model_tf_version} != {tf_version}." - ) - - self.use_keras_api = ( - tf_version.major > 1 - or self.weight_format == KerasModelAdapter.weight_format - ) - - # TODO tf device management - if devices is not None: - logger.warning( - f"Device management is not implemented for tensorflow yet, ignoring the devices {devices}" - ) - - weight_file = self.require_unzipped(weights.source) - self._network = self._get_network(weight_file) - self._internal_output_axes = [ - tuple(a.id for a in get_axes_infos(out)) - for out in model_description.outputs - ] - - # TODO: check how to load tf weights without unzipping - def require_unzipped(self, weight_file: FileSource): - local_weights_file = download(weight_file).path - if isinstance(local_weights_file, ZipPath): - # weights file is in a bioimageio zip package - out_path = ( - Path("bioimageio_unzipped_tf_weights") / local_weights_file.filename - ) - with local_weights_file.open("rb") as src, out_path.open("wb") as dst: - assert not isinstance(src, TextIOWrapper) - copyfileobj(src, dst) - - local_weights_file = out_path - - if zipfile.is_zipfile(local_weights_file): - # weights file itself is a zipfile - out_path = local_weights_file.with_suffix(".unzipped") - with zipfile.ZipFile(local_weights_file, "r") as f: - f.extractall(out_path) - - return out_path - else: - return local_weights_file - - def _get_network( # pyright: ignore[reportUnknownParameterType] - self, weight_file: FileSource - ): - weight_file = self.require_unzipped(weight_file) - assert tf is not None - if self.use_keras_api: - 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) - - # 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( # pyright: ignore[reportUnknownParameterType] - self, *input_tensors: Optional[Tensor] - ): - assert tf is not None - input_keys = [ - ipt.name if isinstance(ipt, v0_4.InputTensorDescr) else ipt.id - for ipt in self.model_description.inputs - ] - output_keys = [ - out.name if isinstance(out, v0_4.OutputTensorDescr) else out.id - for out in self.model_description.outputs - ] - # TODO read from spec - tag = ( # pyright: ignore[reportUnknownVariableType] - tf.saved_model.tag_constants.SERVING - ) - signature_key = ( # pyright: ignore[reportUnknownVariableType] - tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY - ) - - graph = tf.Graph() # pyright: ignore[reportUnknownVariableType] - with graph.as_default(): - with tf.Session( - graph=graph - ) as sess: # pyright: ignore[reportUnknownVariableType] - # load the model and the signature - graph_def = tf.saved_model.loader.load( # pyright: ignore[reportUnknownVariableType] - sess, [tag], self._network - ) - signature = ( # pyright: ignore[reportUnknownVariableType] - graph_def.signature_def - ) - - # get the tensors into the graph - in_names = [ # pyright: ignore[reportUnknownVariableType] - signature[signature_key].inputs[key].name for key in input_keys - ] - out_names = [ # pyright: ignore[reportUnknownVariableType] - signature[signature_key].outputs[key].name for key in output_keys - ] - in_tensors = [ # pyright: ignore[reportUnknownVariableType] - graph.get_tensor_by_name(name) - for name in in_names # pyright: ignore[reportUnknownVariableType] - ] - out_tensors = [ # pyright: ignore[reportUnknownVariableType] - graph.get_tensor_by_name(name) - for name in out_names # pyright: ignore[reportUnknownVariableType] - ] - - # run prediction - res = sess.run( # pyright: ignore[reportUnknownVariableType] - dict( - zip( - out_names, # pyright: ignore[reportUnknownArgumentType] - out_tensors, # pyright: ignore[reportUnknownArgumentType] - ) - ), - dict( - zip( - in_tensors, # pyright: ignore[reportUnknownArgumentType] - input_tensors, - ) - ), - ) - # from dict to list of tensors - res = [ # pyright: ignore[reportUnknownVariableType] - res[out] - for out in out_names # pyright: ignore[reportUnknownVariableType] - ] - - return res # pyright: ignore[reportUnknownVariableType] - - def _forward_keras( # pyright: ignore[reportUnknownParameterType] - self, *input_tensors: Optional[Tensor] - ): - assert self.use_keras_api - assert not isinstance(self._network, str) - assert tf is not None - tf_tensor = [ # pyright: ignore[reportUnknownVariableType] - None if ipt is None else tf.convert_to_tensor(ipt) for ipt in input_tensors - ] - - result = self._network(*tf_tensor) # 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 r.numpy()) - for r in result # pyright: ignore[reportUnknownVariableType] - ] - - def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: - data = [None if ipt is None else ipt.data for ipt in input_tensors] - if self.use_keras_api: - result = self._forward_keras( # pyright: ignore[reportUnknownVariableType] - *data - ) - else: - result = self._forward_tf( # pyright: ignore[reportUnknownVariableType] - *data - ) - - return [ - None if r is None else Tensor(r, dims=axes) - for r, axes in zip( # pyright: ignore[reportUnknownVariableType] - result, # pyright: ignore[reportUnknownArgumentType] - self._internal_output_axes, - ) - ] - - def unload(self) -> None: - logger.warning( - "Device management is not implemented for keras yet, cannot unload model" - ) - - -class TensorflowModelAdapter(TensorflowModelAdapterBase): - weight_format = "tensorflow_saved_model_bundle" - - def __init__( - self, - *, - model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], - devices: Optional[Sequence[str]] = None, - ): - if model_description.weights.tensorflow_saved_model_bundle is None: - raise ValueError("missing tensorflow_saved_model_bundle weights") - - super().__init__( - devices=devices, - weights=model_description.weights.tensorflow_saved_model_bundle, - model_description=model_description, - ) - - -class KerasModelAdapter(TensorflowModelAdapterBase): - weight_format = "keras_hdf5" - - def __init__( - self, - *, - model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], - devices: Optional[Sequence[str]] = None, - ): - if model_description.weights.keras_hdf5 is None: - raise ValueError("missing keras_hdf5 weights") - - super().__init__( - model_description=model_description, - devices=devices, - weights=model_description.weights.keras_hdf5, - ) diff --git a/bioimageio/core/model_adapters/_torchscript_model_adapter.py b/bioimageio/core/model_adapters/_torchscript_model_adapter.py deleted file mode 100644 index 346718a9..00000000 --- a/bioimageio/core/model_adapters/_torchscript_model_adapter.py +++ /dev/null @@ -1,79 +0,0 @@ -import gc -import warnings -from typing import Any, List, Optional, Sequence, Union - -import torch - -from bioimageio.spec._internal.type_guards import is_list, is_ndarray, is_tuple -from bioimageio.spec.model import v0_4, v0_5 -from bioimageio.spec.utils import download - -from ..digest_spec import get_axes_infos -from ..tensor import Tensor -from ._model_adapter import ModelAdapter - - -class TorchscriptModelAdapter(ModelAdapter): - def __init__( - self, - *, - model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], - devices: Optional[Sequence[str]] = None, - ): - super().__init__() - if model_description.weights.torchscript is None: - raise ValueError( - f"No torchscript weights found for model {model_description.name}" - ) - - weight_path = download(model_description.weights.torchscript.source).path - 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._model = self._model.eval() - self._internal_output_axes = [ - tuple(a.id for a in get_axes_infos(out)) - for out in model_description.outputs - ] - - def forward(self, *batch: Optional[Tensor]) -> List[Optional[Tensor]]: - with torch.no_grad(): - torch_tensor = [ - None if b is None else torch.from_numpy(b.data.data).to(self.devices[0]) - for b in batch - ] - _result: Any = self._model.forward(*torch_tensor) - if is_list(_result) or is_tuple(_result): - result: Sequence[Any] = _result - else: - result = [_result] - - result = [ - ( - None - if r is None - else r.cpu().numpy() if isinstance(r, torch.Tensor) else r - ) - for r in result - ] - - assert len(result) == len(self._internal_output_axes) - return [ - None if r is None else Tensor(r, dims=axes) if is_ndarray(r) else r - 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 diff --git a/bioimageio/core/weight_converters/__init__.py b/bioimageio/core/weight_converters/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bioimageio/core/weight_converters/_add_weights.py b/bioimageio/core/weight_converters/_add_weights.py new file mode 100644 index 00000000..76041550 --- /dev/null +++ b/bioimageio/core/weight_converters/_add_weights.py @@ -0,0 +1,25 @@ +from abc import ABC +from typing import Optional, Sequence, Union, assert_never, final + +from bioimageio.spec.model import v0_4, v0_5 + + +def increase_available_weight_formats( + model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], + source_format: v0_5.WeightsFormat, + target_format: v0_5.WeightsFormat, + *, + devices: Optional[Sequence[str]] = None, +): + if not isinstance(model_description, (v0_4.ModelDescr, v0_5.ModelDescr)): + raise TypeError( + f"expected v0_4.ModelDescr or v0_5.ModelDescr, but got {type(model_description)}" + ) + + if (source_format, target_format) == ("pytorch_state_dict", "onnx"): + from .pytorch_to_onnx import convert_pytorch_to_onnx + + else: + raise NotImplementedError( + f"Converting from '{source_format}' to '{target_format}' is not yet implemented. Please create an issue at https://github.com/bioimage-io/core-bioimage-io-python/issues/new/choose" + ) diff --git a/bioimageio/core/weight_converters/pytorch_to_onnx.py b/bioimageio/core/weight_converters/pytorch_to_onnx.py new file mode 100644 index 00000000..acb621e2 --- /dev/null +++ b/bioimageio/core/weight_converters/pytorch_to_onnx.py @@ -0,0 +1,124 @@ +import abc +import os +import shutil +from pathlib import Path +from typing import Any, List, Sequence, Union, cast, no_type_check +from zipfile import ZipFile + +import numpy as np +import torch +from numpy.testing import assert_array_almost_equal +from torch.jit import ScriptModule +from typing_extensions import assert_never + +from bioimageio.core.digest_spec import get_member_id, get_test_inputs +from bioimageio.core.model_adapters._pytorch_model_adapter import PytorchModelAdapter +from bioimageio.spec._internal.io_utils import download +from bioimageio.spec._internal.version_type import Version +from bioimageio.spec.model import v0_4, v0_5 +from bioimageio.spec.model.v0_5 import WeightsEntryDescrBase + + +def convert( + model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], + *, + # output_path: Path, + use_tracing: bool = True, + test_decimal: int = 4, + verbose: bool = False, + opset_version: int = 15, +) -> v0_5.OnnxWeightsDescr: + """ + Convert model weights from the PyTorch state_dict format to the ONNX format. + + # TODO: update Args + Args: + model_descr (Union[v0_4.ModelDescr, v0_5.ModelDescr]): + The model description object that contains the model and its weights. + output_path (Path): + The file path where the ONNX model will be saved. + use_tracing (bool, optional): + Whether to use tracing or scripting to export the ONNX format. Defaults to True. + test_decimal (int, optional): + The decimal precision for comparing the results between the original and converted models. + This is used in the `assert_array_almost_equal` function to check if the outputs match. + Defaults to 4. + verbose (bool, optional): + If True, will print out detailed information during the ONNX export process. Defaults to False. + opset_version (int, optional): + The ONNX opset version to use for the export. Defaults to 15. + Raises: + ValueError: + If the provided model does not have weights in the PyTorch state_dict format. + ImportError: + If ONNX Runtime is not available for checking the exported ONNX model. + ValueError: + If the results before and after weights conversion do not agree. + Returns: + v0_5.OnnxWeightsDescr: + A descriptor object that contains information about the exported ONNX weights. + """ + + state_dict_weights_descr = model_descr.weights.pytorch_state_dict + if state_dict_weights_descr is None: + raise ValueError( + "The provided model does not have weights in the pytorch state dict format" + ) + + with torch.no_grad(): + sample = get_test_inputs(model_descr) + input_data = [ + sample.members[get_member_id(ipt)].data.data for ipt in model_descr.inputs + ] + input_tensors = [torch.from_numpy(ipt) for ipt in input_data] + model = load_torch_model(state_dict_weights_descr) + + expected_tensors = model(*input_tensors) + if isinstance(expected_tensors, torch.Tensor): + expected_tensors = [expected_tensors] + expected_outputs: List[np.ndarray[Any, Any]] = [ + out.numpy() for out in expected_tensors + ] + + if use_tracing: + torch.onnx.export( + model, + (tuple(input_tensors) if len(input_tensors) > 1 else input_tensors[0]), + str(output_path), + verbose=verbose, + opset_version=opset_version, + ) + else: + raise NotImplementedError + + try: + import onnxruntime as rt # pyright: ignore [reportMissingTypeStubs] + except ImportError: + 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)) + onnx_input_node_args = cast( + List[Any], sess.get_inputs() + ) # fixme: remove cast, try using rt.NodeArg instead of Any + onnx_inputs = { + input_name.name: inp + for input_name, inp in zip(onnx_input_node_args, input_data) + } + outputs = cast( + Sequence[np.ndarray[Any, Any]], sess.run(None, onnx_inputs) + ) # FIXME: remove cast + + try: + for exp, out in zip(expected_outputs, outputs): + assert_array_almost_equal(exp, out, decimal=test_decimal) + except AssertionError as e: + 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 + ) From 7fea808f0570cfec4b74f36efb2835ab18bc3e9c Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 5 Dec 2024 11:39:24 +0100 Subject: [PATCH 17/47] add summary_path arg --- bioimageio/core/cli.py | 4 ++++ bioimageio/core/commands.py | 6 ++++++ 2 files changed, 10 insertions(+) diff --git a/bioimageio/core/cli.py b/bioimageio/core/cli.py index fad44ab3..b81b1e5e 100644 --- a/bioimageio/core/cli.py +++ b/bioimageio/core/cli.py @@ -133,6 +133,9 @@ class TestCmd(CmdBase, WithSource): decimal: int = 4 """Precision for numerical comparisons""" + summary_path: Optional[Path] = None + """Path to save validation summary as JSON file.""" + def run(self): sys.exit( test( @@ -140,6 +143,7 @@ def run(self): weight_format=self.weight_format, devices=self.devices, decimal=self.decimal, + summary_path=self.summary_path, ) ) diff --git a/bioimageio/core/commands.py b/bioimageio/core/commands.py index c71d495f..6ad54ab7 100644 --- a/bioimageio/core/commands.py +++ b/bioimageio/core/commands.py @@ -1,6 +1,7 @@ """These functions implement the logic of the bioimageio command line interface defined in `bioimageio.core.cli`.""" +import json from pathlib import Path from typing import Optional, Sequence, Union @@ -26,6 +27,7 @@ def test( weight_format: WeightFormatArgAll = "all", devices: Optional[Union[str, Sequence[str]]] = None, decimal: int = 4, + summary_path: Optional[Path] = None, ) -> int: """test a bioimageio resource @@ -35,6 +37,7 @@ def test( weight_format: (model only) The weight format to use devices: Device(s) to use for testing decimal: Precision for numerical comparisons + summary_path: Path to save validation summary as JSON file. """ if isinstance(descr, InvalidDescr): descr.validation_summary.display() @@ -47,6 +50,9 @@ def test( decimal=decimal, ) summary.display() + if summary_path is not None: + _ = summary_path.write_text(summary.model_dump_json(indent=4)) + return 0 if summary.status == "passed" else 1 From 4564f7ca92085b81adbce41d94530ab95f9276da Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 5 Dec 2024 21:20:24 +0100 Subject: [PATCH 18/47] update annotation --- bioimageio/core/_resource_tests.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bioimageio/core/_resource_tests.py b/bioimageio/core/_resource_tests.py index 6ace6d5c..23a9b8f8 100644 --- a/bioimageio/core/_resource_tests.py +++ b/bioimageio/core/_resource_tests.py @@ -94,7 +94,7 @@ def enable_determinism(mode: Literal["seed_only", "full"]): def test_model( - source: Union[v0_5.ModelDescr, PermissiveFileSource], + source: Union[v0_4.ModelDescr, v0_5.ModelDescr, PermissiveFileSource], weight_format: Optional[WeightsFormat] = None, devices: Optional[List[str]] = None, absolute_tolerance: float = 1.5e-4, From 00e6ba15236cc2182a2dedd1c5cb783064661239 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 5 Dec 2024 21:20:53 +0100 Subject: [PATCH 19/47] fix tf seeding --- bioimageio/core/_resource_tests.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bioimageio/core/_resource_tests.py b/bioimageio/core/_resource_tests.py index 23a9b8f8..88323492 100644 --- a/bioimageio/core/_resource_tests.py +++ b/bioimageio/core/_resource_tests.py @@ -81,11 +81,11 @@ def enable_determinism(mode: Literal["seed_only", "full"]): try: try: - import tensorflow as tf # pyright: ignore[reportMissingImports] + import tensorflow as tf except ImportError: pass else: - tf.random.seed(0) + tf.random.set_seed(0) if mode == "full": tf.config.experimental.enable_op_determinism() # TODO: find possibility to switch it off again?? From 5d1e2ce1543762d4a0aa3d3bfc3a7512d039f63c Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 5 Dec 2024 22:06:12 +0100 Subject: [PATCH 20/47] expose test_description_in_conda_env --- README.md | 5 + bioimageio/core/__init__.py | 2 + bioimageio/core/_dynamic_conda_env.py | 165 ++++++++++++++++++++++++++ 3 files changed, 172 insertions(+) create mode 100644 bioimageio/core/_dynamic_conda_env.py diff --git a/README.md b/README.md index 49e6cbbd..2233ec97 100644 --- a/README.md +++ b/README.md @@ -375,6 +375,11 @@ The model specification and its validation tools can be found at ValidationSummary: + """Run test_model in a dedicated conda env + + Args: + source: Path or URL to model description. + weight_format: Weight format to test. + Default: All weight formats present in **source**. + conda_env: conda environment including bioimageio.core dependency. + Default: Use `bioimageio.spec.get_conda_env` to obtain a model weight + specific conda environment. + devices: Devices to test with, e.g. 'cpu', 'cuda'. + Default (may be weight format dependent): ['cuda'] if available, ['cpu'] otherwise. + absolute_tolerance: Maximum absolute tolerance of reproduced output tensors. + relative_tolerance: Maximum relative tolerance of reproduced output tensors. + determinism: Modes to improve reproducibility of test outputs. + run_command: Function to execute terminal commands. + """ + + try: + run_command(["which", "conda"]) + except Exception as e: + raise RuntimeError("Conda not available") from e + + descr = load_description(source) + if not isinstance(descr, (v0_4.ModelDescr, v0_5.ModelDescr)): + raise NotImplementedError("Not yet implemented for non-model resources") + + if weight_format is None: + all_present_wfs = [ + wf for wf in get_args(WeightsFormat) if getattr(descr.weights, wf) + ] + ignore_wfs = [wf for wf in all_present_wfs if wf in ["tensorflow_js"]] + logger.info( + "Found weight formats {}. Start testing all{}...", + all_present_wfs, + f" (except: {', '.join(ignore_wfs)}) " if ignore_wfs else "", + ) + summary = test_description_in_env( + source, + weight_format=all_present_wfs[0], + devices=devices, + absolute_tolerance=absolute_tolerance, + relative_tolerance=relative_tolerance, + determinism=determinism, + ) + for wf in all_present_wfs[1:]: + additional_summary = test_description_in_env( + source, + weight_format=all_present_wfs[0], + devices=devices, + absolute_tolerance=absolute_tolerance, + relative_tolerance=relative_tolerance, + determinism=determinism, + ) + for d in additional_summary.details: + # TODO: filter reduntant details; group details + summary.add_detail(d) + return summary + + if weight_format == "pytorch_state_dict": + wf = descr.weights.pytorch_state_dict + elif weight_format == "torchscript": + wf = descr.weights.torchscript + elif weight_format == "keras_hdf5": + wf = descr.weights.keras_hdf5 + elif weight_format == "onnx": + wf = descr.weights.onnx + elif weight_format == "tensorflow_saved_model_bundle": + wf = descr.weights.tensorflow_saved_model_bundle + elif weight_format == "tensorflow_js": + raise RuntimeError( + "testing 'tensorflow_js' is not supported by bioimageio.core" + ) + else: + assert_never(weight_format) + + assert wf is not None + if conda_env is None: + conda_env = get_conda_env(entry=wf) + + # remove name as we crate a name based on the env description hash value + conda_env.name = None + + dumped_env = conda_env.model_dump(mode="json", exclude_none=True) + if not is_yaml_value(dumped_env): + raise ValueError(f"Failed to dump conda env to valid YAML {conda_env}") + + env_io = StringIO() + write_yaml(dumped_env, file=env_io) + encoded_env = env_io.getvalue().encode() + env_name = sha256(encoded_env).hexdigest() + + with TemporaryDirectory() as _d: + folder = Path(_d) + try: + run_command(["conda", "activate", env_name]) + except Exception: + path = folder / "env.yaml" + _ = path.write_bytes(encoded_env) + + run_command( + ["conda", "env", "create", "--file", str(path), "--name", env_name] + ) + run_command(["conda", "activate", env_name]) + + summary_path = folder / "summary.json" + run_command( + [ + "conda", + "run", + "-n", + env_name, + "bioimageio", + "test", + str(source), + "--summary-path", + str(summary_path), + ] + ) + return ValidationSummary.model_validate_json(summary_path.read_bytes()) From df36d15e413cebe0be3b1b6919a4788b0a0a5ab2 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 5 Dec 2024 22:06:26 +0100 Subject: [PATCH 21/47] docstring formatting --- bioimageio/core/cli.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bioimageio/core/cli.py b/bioimageio/core/cli.py index b81b1e5e..ad75a51f 100644 --- a/bioimageio/core/cli.py +++ b/bioimageio/core/cli.py @@ -113,14 +113,14 @@ def descr_id(self) -> str: class ValidateFormatCmd(CmdBase, WithSource): - """validate the meta data format of a bioimageio resource.""" + """Validate the meta data format of a bioimageio resource.""" def run(self): sys.exit(validate_format(self.descr)) class TestCmd(CmdBase, WithSource): - """Test a bioimageio resource (beyond meta data formatting)""" + """Test a bioimageio resource (beyond meta data formatting).""" weight_format: WeightFormatArgAll = "all" """The weight format to limit testing to. @@ -149,7 +149,7 @@ def run(self): class PackageCmd(CmdBase, WithSource): - """save a resource's metadata with its associated files.""" + """Save a resource's metadata with its associated files.""" path: CliPositionalArg[Path] """The path to write the (zipped) package to. From 376507f24972356258c65a9f2509d68e0fecd567 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 6 Dec 2024 10:54:51 +0100 Subject: [PATCH 22/47] absorb test_description_in_conda_env into test_description --- bioimageio/core/__init__.py | 1 - bioimageio/core/_dynamic_conda_env.py | 165 ------------------- bioimageio/core/_resource_tests.py | 228 ++++++++++++++++++++++++-- bioimageio/core/cli.py | 18 +- bioimageio/core/commands.py | 15 +- 5 files changed, 236 insertions(+), 191 deletions(-) delete mode 100644 bioimageio/core/_dynamic_conda_env.py diff --git a/bioimageio/core/__init__.py b/bioimageio/core/__init__.py index 9e4d83d9..f47f8f63 100644 --- a/bioimageio/core/__init__.py +++ b/bioimageio/core/__init__.py @@ -32,7 +32,6 @@ stat_measures, tensor, ) -from ._dynamic_conda_env import test_description_in_conda_env from ._prediction_pipeline import PredictionPipeline, create_prediction_pipeline from ._resource_tests import ( enable_determinism, diff --git a/bioimageio/core/_dynamic_conda_env.py b/bioimageio/core/_dynamic_conda_env.py deleted file mode 100644 index 26e3ba35..00000000 --- a/bioimageio/core/_dynamic_conda_env.py +++ /dev/null @@ -1,165 +0,0 @@ -import subprocess -from hashlib import sha256 -from io import StringIO -from pathlib import Path -from tempfile import TemporaryDirectory -from typing import ( - Callable, - List, - Literal, - Optional, - Sequence, - assert_never, -) - -from loguru import logger -from typing_extensions import get_args - -from bioimageio.spec import ( - BioimageioCondaEnv, - ValidationSummary, - get_conda_env, - load_description, -) -from bioimageio.spec._internal.io import is_yaml_value -from bioimageio.spec._internal.io_utils import write_yaml -from bioimageio.spec.common import PermissiveFileSource -from bioimageio.spec.model import v0_4, v0_5 -from bioimageio.spec.model.v0_5 import WeightsFormat - - -def default_run_command(args: Sequence[str]): - logger.info("running '{}'...", " ".join(args)) - _ = subprocess.run(args, shell=True, text=True, check=True) - - -def test_description_in_conda_env( - source: PermissiveFileSource, - *, - weight_format: Optional[WeightsFormat] = None, - conda_env: Optional[BioimageioCondaEnv] = None, - devices: Optional[List[str]] = None, - absolute_tolerance: float = 1.5e-4, - relative_tolerance: float = 1e-4, - determinism: Literal["seed_only", "full"] = "seed_only", - run_command: Callable[[Sequence[str]], None] = default_run_command, -) -> ValidationSummary: - """Run test_model in a dedicated conda env - - Args: - source: Path or URL to model description. - weight_format: Weight format to test. - Default: All weight formats present in **source**. - conda_env: conda environment including bioimageio.core dependency. - Default: Use `bioimageio.spec.get_conda_env` to obtain a model weight - specific conda environment. - devices: Devices to test with, e.g. 'cpu', 'cuda'. - Default (may be weight format dependent): ['cuda'] if available, ['cpu'] otherwise. - absolute_tolerance: Maximum absolute tolerance of reproduced output tensors. - relative_tolerance: Maximum relative tolerance of reproduced output tensors. - determinism: Modes to improve reproducibility of test outputs. - run_command: Function to execute terminal commands. - """ - - try: - run_command(["which", "conda"]) - except Exception as e: - raise RuntimeError("Conda not available") from e - - descr = load_description(source) - if not isinstance(descr, (v0_4.ModelDescr, v0_5.ModelDescr)): - raise NotImplementedError("Not yet implemented for non-model resources") - - if weight_format is None: - all_present_wfs = [ - wf for wf in get_args(WeightsFormat) if getattr(descr.weights, wf) - ] - ignore_wfs = [wf for wf in all_present_wfs if wf in ["tensorflow_js"]] - logger.info( - "Found weight formats {}. Start testing all{}...", - all_present_wfs, - f" (except: {', '.join(ignore_wfs)}) " if ignore_wfs else "", - ) - summary = test_description_in_env( - source, - weight_format=all_present_wfs[0], - devices=devices, - absolute_tolerance=absolute_tolerance, - relative_tolerance=relative_tolerance, - determinism=determinism, - ) - for wf in all_present_wfs[1:]: - additional_summary = test_description_in_env( - source, - weight_format=all_present_wfs[0], - devices=devices, - absolute_tolerance=absolute_tolerance, - relative_tolerance=relative_tolerance, - determinism=determinism, - ) - for d in additional_summary.details: - # TODO: filter reduntant details; group details - summary.add_detail(d) - return summary - - if weight_format == "pytorch_state_dict": - wf = descr.weights.pytorch_state_dict - elif weight_format == "torchscript": - wf = descr.weights.torchscript - elif weight_format == "keras_hdf5": - wf = descr.weights.keras_hdf5 - elif weight_format == "onnx": - wf = descr.weights.onnx - elif weight_format == "tensorflow_saved_model_bundle": - wf = descr.weights.tensorflow_saved_model_bundle - elif weight_format == "tensorflow_js": - raise RuntimeError( - "testing 'tensorflow_js' is not supported by bioimageio.core" - ) - else: - assert_never(weight_format) - - assert wf is not None - if conda_env is None: - conda_env = get_conda_env(entry=wf) - - # remove name as we crate a name based on the env description hash value - conda_env.name = None - - dumped_env = conda_env.model_dump(mode="json", exclude_none=True) - if not is_yaml_value(dumped_env): - raise ValueError(f"Failed to dump conda env to valid YAML {conda_env}") - - env_io = StringIO() - write_yaml(dumped_env, file=env_io) - encoded_env = env_io.getvalue().encode() - env_name = sha256(encoded_env).hexdigest() - - with TemporaryDirectory() as _d: - folder = Path(_d) - try: - run_command(["conda", "activate", env_name]) - except Exception: - path = folder / "env.yaml" - _ = path.write_bytes(encoded_env) - - run_command( - ["conda", "env", "create", "--file", str(path), "--name", env_name] - ) - run_command(["conda", "activate", env_name]) - - summary_path = folder / "summary.json" - run_command( - [ - "conda", - "run", - "-n", - env_name, - "bioimageio", - "test", - str(source), - "--summary-path", - str(summary_path), - ] - ) - return ValidationSummary.model_validate_json(summary_path.read_bytes()) diff --git a/bioimageio/core/_resource_tests.py b/bioimageio/core/_resource_tests.py index 88323492..bcffbc53 100644 --- a/bioimageio/core/_resource_tests.py +++ b/bioimageio/core/_resource_tests.py @@ -1,21 +1,43 @@ +import hashlib +import platform +import subprocess import traceback import warnings +from io import StringIO from itertools import product -from typing import Dict, Hashable, List, Literal, Optional, Sequence, Set, Tuple, Union +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import ( + Callable, + Dict, + Hashable, + List, + Literal, + Optional, + Sequence, + Set, + Tuple, + Union, +) import numpy as np from loguru import logger +from typing_extensions import assert_never, get_args from bioimageio.spec import ( + BioimageioCondaEnv, InvalidDescr, ResourceDescr, build_description, dump_description, + get_conda_env, load_description, + save_bioimageio_package, ) from bioimageio.spec._internal.common_nodes import ResourceDescrBase +from bioimageio.spec._internal.io import is_yaml_value +from bioimageio.spec._internal.io_utils import read_yaml, write_yaml from bioimageio.spec.common import BioimageioYamlContent, PermissiveFileSource -from bioimageio.spec.get_conda_env import get_conda_env from bioimageio.spec.model import v0_4, v0_5 from bioimageio.spec.model.v0_5 import WeightsFormat from bioimageio.spec.summary import ( @@ -116,6 +138,11 @@ def test_model( ) +def default_run_command(args: Sequence[str]): + logger.info("running '{}'...", " ".join(args)) + _ = subprocess.run(args, shell=True, text=True, check=True) + + def test_description( source: Union[ResourceDescr, PermissiveFileSource, BioimageioYamlContent], *, @@ -127,20 +154,193 @@ def test_description( decimal: Optional[int] = None, determinism: Literal["seed_only", "full"] = "seed_only", expected_type: Optional[str] = None, + runtime_env: Union[ + Literal["currently-active", "as-described"], Path, BioimageioCondaEnv + ] = ("currently-active"), + run_command: Callable[[Sequence[str]], None] = default_run_command, ) -> ValidationSummary: - """Test a bioimage.io resource dynamically, e.g. prediction of test tensors for models""" - rd = load_description_and_test( - source, - format_version=format_version, - weight_format=weight_format, - devices=devices, - absolute_tolerance=absolute_tolerance, - relative_tolerance=relative_tolerance, - decimal=decimal, - determinism=determinism, - expected_type=expected_type, + """Test a bioimage.io resource dynamically, e.g. prediction of test tensors for models. + + Args: + source: model description source. + weight_format: Weight format to test. + Default: All weight formats present in **source**. + devices: Devices to test with, e.g. 'cpu', 'cuda'. + Default (may be weight format dependent): ['cuda'] if available, ['cpu'] otherwise. + absolute_tolerance: Maximum absolute tolerance of reproduced output tensors. + relative_tolerance: Maximum relative tolerance of reproduced output tensors. + determinism: Modes to improve reproducibility of test outputs. + runtime_env: (Experimental feature!) The Python environment to run the tests in + - `"currently-active"`: Use active Python interpreter. + - `"as-described"`: Use `bioimageio.spec.get_conda_env` to generate a conda + environment YAML file based on the model weights description. + - A `BioimageioCondaEnv` or a path to a conda environment YAML file. + Note: The `bioimageio.core` dependency will be added automatically if not present. + run_command: (Experimental feature!) Function to execute (conda) terminal commands in a subprocess + (ignored if **runtime_env** is `"currently-active"`). + """ + if runtime_env == "currently-active": + rd = load_description_and_test( + source, + format_version=format_version, + weight_format=weight_format, + devices=devices, + absolute_tolerance=absolute_tolerance, + relative_tolerance=relative_tolerance, + decimal=decimal, + determinism=determinism, + expected_type=expected_type, + ) + return rd.validation_summary + + if runtime_env == "as-described": + conda_env = None + elif isinstance(runtime_env, (str, Path)): + conda_env = BioimageioCondaEnv.model_validate(read_yaml(Path(runtime_env))) + elif isinstance(runtime_env, BioimageioCondaEnv): + conda_env = runtime_env + else: + assert_never(runtime_env) + + with TemporaryDirectory(ignore_cleanup_errors=True) as _d: + working_dir = Path(_d) + if isinstance(source, (dict, ResourceDescrBase)): + file_source = save_bioimageio_package( + source, output_path=working_dir / "package.zip" + ) + else: + file_source = source + + return _test_in_env( + file_source, + working_dir=working_dir, + weight_format=weight_format, + conda_env=conda_env, + devices=devices, + absolute_tolerance=absolute_tolerance, + relative_tolerance=relative_tolerance, + determinism=determinism, + run_command=run_command, + ) + + +def _test_in_env( + source: PermissiveFileSource, + *, + working_dir: Path, + weight_format: Optional[WeightsFormat], + conda_env: Optional[BioimageioCondaEnv], + devices: Optional[Sequence[str]], + absolute_tolerance: float, + relative_tolerance: float, + determinism: Literal["seed_only", "full"], + run_command: Callable[[Sequence[str]], None], +) -> ValidationSummary: + descr = load_description(source) + + if not isinstance(descr, (v0_4.ModelDescr, v0_5.ModelDescr)): + raise NotImplementedError("Not yet implemented for non-model resources") + + if weight_format is None: + all_present_wfs = [ + wf for wf in get_args(WeightsFormat) if getattr(descr.weights, wf) + ] + ignore_wfs = [wf for wf in all_present_wfs if wf in ["tensorflow_js"]] + logger.info( + "Found weight formats {}. Start testing all{}...", + all_present_wfs, + f" (except: {', '.join(ignore_wfs)}) " if ignore_wfs else "", + ) + summary = _test_in_env( + source, + working_dir=working_dir / all_present_wfs[0], + weight_format=all_present_wfs[0], + devices=devices, + absolute_tolerance=absolute_tolerance, + relative_tolerance=relative_tolerance, + determinism=determinism, + conda_env=conda_env, + run_command=run_command, + ) + for wf in all_present_wfs[1:]: + additional_summary = _test_in_env( + source, + working_dir=working_dir / wf, + weight_format=wf, + devices=devices, + absolute_tolerance=absolute_tolerance, + relative_tolerance=relative_tolerance, + determinism=determinism, + conda_env=conda_env, + run_command=run_command, + ) + for d in additional_summary.details: + # TODO: filter reduntant details; group details + summary.add_detail(d) + return summary + + if weight_format == "pytorch_state_dict": + wf = descr.weights.pytorch_state_dict + elif weight_format == "torchscript": + wf = descr.weights.torchscript + elif weight_format == "keras_hdf5": + wf = descr.weights.keras_hdf5 + elif weight_format == "onnx": + wf = descr.weights.onnx + elif weight_format == "tensorflow_saved_model_bundle": + wf = descr.weights.tensorflow_saved_model_bundle + elif weight_format == "tensorflow_js": + raise RuntimeError( + "testing 'tensorflow_js' is not supported by bioimageio.core" + ) + else: + assert_never(weight_format) + + assert wf is not None + if conda_env is None: + conda_env = get_conda_env(entry=wf) + + # remove name as we crate a name based on the env description hash value + conda_env.name = None + + dumped_env = conda_env.model_dump(mode="json", exclude_none=True) + if not is_yaml_value(dumped_env): + raise ValueError(f"Failed to dump conda env to valid YAML {conda_env}") + + env_io = StringIO() + write_yaml(dumped_env, file=env_io) + encoded_env = env_io.getvalue().encode() + env_name = hashlib.sha256(encoded_env).hexdigest() + + try: + run_command(["where" if platform.system() == "Windows" else "which", "conda"]) + except Exception as e: + raise RuntimeError("Conda not available") from e + + working_dir.mkdir(parents=True, exist_ok=True) + try: + run_command(["conda", "activate", env_name]) + except Exception: + path = working_dir / "env.yaml" + _ = path.write_bytes(encoded_env) + logger.debug("written conda env to {}", path) + run_command(["conda", "env", "create", f"--file={path}", f"--name={env_name}"]) + run_command(["conda", "activate", env_name]) + + summary_path = working_dir / "summary.json" + run_command( + [ + "conda", + "run", + "-n", + env_name, + "bioimageio", + "test", + str(source), + f"--summary-path={summary_path}", + ] ) - return rd.validation_summary + return ValidationSummary.model_validate_json(summary_path.read_bytes()) def load_description_and_test( diff --git a/bioimageio/core/cli.py b/bioimageio/core/cli.py index ad75a51f..1fc95310 100644 --- a/bioimageio/core/cli.py +++ b/bioimageio/core/cli.py @@ -18,6 +18,7 @@ Dict, Iterable, List, + Literal, Mapping, Optional, Sequence, @@ -133,7 +134,19 @@ class TestCmd(CmdBase, WithSource): decimal: int = 4 """Precision for numerical comparisons""" - summary_path: Optional[Path] = None + runtime_env: Union[Literal["currently-active", "as-described"], Path] = Field( + "currently-active", alias="runtime-env" + ) + """The python environment to run the tests in + + - `"currently-active"`: use active Python interpreter + - `"as-described"`: generate a conda environment YAML file based on the model + weights description. + - A path to a conda environment YAML. + Note: The `bioimageio.core` dependency will be added automatically if not present. + """ + + summary_path: Optional[Path] = Field(None, alias="summary-path") """Path to save validation summary as JSON file.""" def run(self): @@ -144,6 +157,7 @@ def run(self): devices=self.devices, decimal=self.decimal, summary_path=self.summary_path, + runtime_env=self.runtime_env, ) ) @@ -555,10 +569,10 @@ def input_dataset(stat: Stat): class Bioimageio( BaseSettings, + cli_implicit_flags=True, cli_parse_args=True, cli_prog_name="bioimageio", cli_use_class_docs_for_groups=True, - cli_implicit_flags=True, use_attribute_docstrings=True, ): """bioimageio - CLI for bioimage.io resources 🦒""" diff --git a/bioimageio/core/commands.py b/bioimageio/core/commands.py index 6ad54ab7..9804a93e 100644 --- a/bioimageio/core/commands.py +++ b/bioimageio/core/commands.py @@ -1,7 +1,6 @@ """These functions implement the logic of the bioimageio command line interface defined in `bioimageio.core.cli`.""" -import json from pathlib import Path from typing import Optional, Sequence, Union @@ -28,16 +27,13 @@ def test( devices: Optional[Union[str, Sequence[str]]] = None, decimal: int = 4, summary_path: Optional[Path] = None, + runtime_env: Union[ + Literal["currently-active", "as-described"], Path + ] = "currently-active", ) -> int: - """test a bioimageio resource + """Test a bioimageio resource. - Args: - source: Path or URL to the bioimageio resource description file - (bioimageio.yaml or rdf.yaml) or to a zipped resource - weight_format: (model only) The weight format to use - devices: Device(s) to use for testing - decimal: Precision for numerical comparisons - summary_path: Path to save validation summary as JSON file. + Arguments as described in `bioimageio.core.cli.TestCmd` """ if isinstance(descr, InvalidDescr): descr.validation_summary.display() @@ -48,6 +44,7 @@ def test( weight_format=None if weight_format == "all" else weight_format, devices=[devices] if isinstance(devices, str) else devices, decimal=decimal, + runtime_env=runtime_env, ) summary.display() if summary_path is not None: From 96905740bf1fdacd4bcf4459d4aa9c067a76ad97 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 6 Dec 2024 15:56:45 +0100 Subject: [PATCH 23/47] all model adapters in backends --- bioimageio/core/__init__.py | 3 +- bioimageio/core/_create_model_adapter.py | 127 -------- bioimageio/core/_model_adapter.py | 93 ------ bioimageio/core/_resource_tests.py | 3 +- bioimageio/core/backend/__init__.py | 0 bioimageio/core/backends/__init__.py | 3 + .../_model_adapter.py | 36 +-- .../keras.py => backends/keras_backend.py} | 41 ++- bioimageio/core/backends/onnx_backend.py | 60 ++++ .../pytorch_backend.py} | 0 .../core/backends/tensorflow_backend.py | 289 ++++++++++++++++++ .../core/backends/torchscript_backend.py | 79 +++++ bioimageio/core/model_adapters.py | 18 +- .../model_adapters/_pytorch_model_adapter.py | 0 14 files changed, 482 insertions(+), 270 deletions(-) delete mode 100644 bioimageio/core/_create_model_adapter.py delete mode 100644 bioimageio/core/_model_adapter.py delete mode 100644 bioimageio/core/backend/__init__.py create mode 100644 bioimageio/core/backends/__init__.py rename bioimageio/core/{model_adapters => backends}/_model_adapter.py (84%) rename bioimageio/core/{backend/keras.py => backends/keras_backend.py} (74%) create mode 100644 bioimageio/core/backends/onnx_backend.py rename bioimageio/core/{backend/pytorch.py => backends/pytorch_backend.py} (100%) create mode 100644 bioimageio/core/backends/tensorflow_backend.py create mode 100644 bioimageio/core/backends/torchscript_backend.py delete mode 100644 bioimageio/core/model_adapters/_pytorch_model_adapter.py diff --git a/bioimageio/core/__init__.py b/bioimageio/core/__init__.py index f47f8f63..a8dd1043 100644 --- a/bioimageio/core/__init__.py +++ b/bioimageio/core/__init__.py @@ -41,6 +41,7 @@ ) from ._settings import settings from .axis import Axis, AxisId +from .backends import create_model_adapter from .block_meta import BlockMeta from .common import MemberId from .prediction import predict, predict_many @@ -73,6 +74,7 @@ "commands", "common", "compute_dataset_measures", + "create_model_adapter", "create_prediction_pipeline", "digest_spec", "dump_description", @@ -104,7 +106,6 @@ "Stat", "tensor", "Tensor", - "test_description_in_conda_env", "test_description", "test_model", "test_resource", diff --git a/bioimageio/core/_create_model_adapter.py b/bioimageio/core/_create_model_adapter.py deleted file mode 100644 index ee79f260..00000000 --- a/bioimageio/core/_create_model_adapter.py +++ /dev/null @@ -1,127 +0,0 @@ -import warnings -from abc import abstractmethod -from typing import List, Optional, Sequence, Tuple, Union, final - -from bioimageio.spec.model import v0_4, v0_5 - -from ._model_adapter import ( - DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER, - ModelAdapter, - WeightsFormat, -) -from .tensor import Tensor - - -def create_model_adapter( - model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], - *, - devices: Optional[Sequence[str]] = None, - weight_format_priority_order: Optional[Sequence[WeightsFormat]] = None, -): - """ - 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 - """ - if not isinstance(model_description, (v0_4.ModelDescr, v0_5.ModelDescr)): - raise TypeError( - f"expected v0_4.ModelDescr or v0_5.ModelDescr, but got {type(model_description)}" - ) - - weights = model_description.weights - errors: List[Tuple[WeightsFormat, Exception]] = [] - weight_format_priority_order = ( - DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER - if weight_format_priority_order is None - else weight_format_priority_order - ) - # limit weight formats to the ones present - weight_format_priority_order = [ - w for w in weight_format_priority_order if getattr(weights, w) is not None - ] - - for wf in weight_format_priority_order: - if wf == "pytorch_state_dict" and weights.pytorch_state_dict is not None: - try: - from .model_adapters_old._pytorch_model_adapter import ( - PytorchModelAdapter, - ) - - return PytorchModelAdapter( - outputs=model_description.outputs, - weights=weights.pytorch_state_dict, - devices=devices, - ) - except Exception as e: - errors.append((wf, e)) - elif ( - wf == "tensorflow_saved_model_bundle" - and weights.tensorflow_saved_model_bundle is not None - ): - try: - from .model_adapters_old._tensorflow_model_adapter import ( - TensorflowModelAdapter, - ) - - return TensorflowModelAdapter( - model_description=model_description, devices=devices - ) - except Exception as e: - errors.append((wf, e)) - elif wf == "onnx" and weights.onnx is not None: - try: - from .model_adapters_old._onnx_model_adapter import ONNXModelAdapter - - return ONNXModelAdapter( - model_description=model_description, devices=devices - ) - except Exception as e: - errors.append((wf, e)) - elif wf == "torchscript" and weights.torchscript is not None: - try: - from .model_adapters_old._torchscript_model_adapter import ( - TorchscriptModelAdapter, - ) - - return TorchscriptModelAdapter( - model_description=model_description, devices=devices - ) - except Exception as e: - errors.append((wf, e)) - elif wf == "keras_hdf5" and weights.keras_hdf5 is not None: - # keras can either be installed as a separate package or used as part of tensorflow - # we try to first import the keras model adapter using the separate package and, - # if it is not available, try to load the one using tf - try: - from .backend.keras import ( - KerasModelAdapter, - keras, # type: ignore - ) - - if keras is None: - from .model_adapters_old._tensorflow_model_adapter import ( - KerasModelAdapter, - ) - - return KerasModelAdapter( - model_description=model_description, devices=devices - ) - except Exception as e: - errors.append((wf, e)) - - assert errors - if len(weight_format_priority_order) == 1: - assert len(errors) == 1 - raise ValueError( - f"The '{weight_format_priority_order[0]}' model adapter could not be created" - + f" in this environment:\n{errors[0][1].__class__.__name__}({errors[0][1]}).\n\n" - ) from errors[0][1] - - else: - error_list = "\n - ".join( - f"{wf}: {e.__class__.__name__}({e})" for wf, e in errors - ) - raise ValueError( - "None of the weight format specific model adapters could be created" - + f" in this environment. Errors are:\n\n{error_list}.\n\n" - ) diff --git a/bioimageio/core/_model_adapter.py b/bioimageio/core/_model_adapter.py deleted file mode 100644 index 0438d35e..00000000 --- a/bioimageio/core/_model_adapter.py +++ /dev/null @@ -1,93 +0,0 @@ -import warnings -from abc import ABC, abstractmethod -from typing import List, Optional, Sequence, Tuple, Union, final - -from bioimageio.spec.model import v0_4, v0_5 - -from .tensor import Tensor - -WeightsFormat = Union[v0_4.WeightsFormat, v0_5.WeightsFormat] - -__all__ = [ - "ModelAdapter", - "create_model_adapter", - "get_weight_formats", -] - -# Known weight formats in order of priority -# First match wins -DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER: Tuple[WeightsFormat, ...] = ( - "pytorch_state_dict", - "tensorflow_saved_model_bundle", - "torchscript", - "onnx", - "keras_hdf5", -) - - -class ModelAdapter(ABC): - """ - Represents model *without* any preprocessing or postprocessing. - - ``` - from bioimageio.core import load_description - - model = load_description(...) - - # option 1: - adapter = ModelAdapter.create(model) - adapter.forward(...) - adapter.unload() - - # option 2: - with ModelAdapter.create(model) as adapter: - adapter.forward(...) - ``` - """ - - @final - @classmethod - def create( - cls, - model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], - *, - devices: Optional[Sequence[str]] = None, - weight_format_priority_order: Optional[Sequence[WeightsFormat]] = None, - ): - """ - 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 - """ - from ._create_model_adapter import create_model_adapter - - return create_model_adapter( - model_description, - devices=devices, - weight_format_priority_order=weight_format_priority_order, - ) - - @final - def load(self, *, devices: Optional[Sequence[str]] = None) -> None: - warnings.warn("Deprecated. ModelAdapter is loaded on initialization") - - @abstractmethod - def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: - """ - Run forward pass of model to get model predictions - """ - # TODO: handle tensor.transpose in here and make _forward_impl the abstract impl - - @abstractmethod - def unload(self): - """ - Unload model from any devices, freeing their memory. - The moder adapter should be considered unusable afterwards. - """ - - -def get_weight_formats() -> List[str]: - """ - Return list of supported weight types - """ - return list(DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER) diff --git a/bioimageio/core/_resource_tests.py b/bioimageio/core/_resource_tests.py index 8f24d363..e6675b73 100644 --- a/bioimageio/core/_resource_tests.py +++ b/bioimageio/core/_resource_tests.py @@ -37,6 +37,7 @@ from bioimageio.spec._internal.common_nodes import ResourceDescrBase from bioimageio.spec._internal.io import is_yaml_value from bioimageio.spec._internal.io_utils import read_yaml, write_yaml +from bioimageio.spec.common import BioimageioYamlContent, PermissiveFileSource, Sha256 from bioimageio.spec.model import v0_4, v0_5 from bioimageio.spec.model.v0_5 import WeightsFormat from bioimageio.spec.summary import ( @@ -192,7 +193,7 @@ def test_description( decimal=decimal, determinism=determinism, expected_type=expected_type, - sha256=sha256, + sha256=sha256, ) return rd.validation_summary diff --git a/bioimageio/core/backend/__init__.py b/bioimageio/core/backend/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/bioimageio/core/backends/__init__.py b/bioimageio/core/backends/__init__.py new file mode 100644 index 00000000..c39b58b5 --- /dev/null +++ b/bioimageio/core/backends/__init__.py @@ -0,0 +1,3 @@ +from ._model_adapter import create_model_adapter + +__all__ = ["create_model_adapter"] diff --git a/bioimageio/core/model_adapters/_model_adapter.py b/bioimageio/core/backends/_model_adapter.py similarity index 84% rename from bioimageio/core/model_adapters/_model_adapter.py rename to bioimageio/core/backends/_model_adapter.py index 3921f81b..66153f09 100644 --- a/bioimageio/core/model_adapters/_model_adapter.py +++ b/bioimageio/core/backends/_model_adapter.py @@ -73,7 +73,7 @@ def create( for wf in weight_format_priority_order: if wf == "pytorch_state_dict" and weights.pytorch_state_dict is not None: try: - from ._pytorch_model_adapter import PytorchModelAdapter + from .pytorch_backend import PytorchModelAdapter return PytorchModelAdapter( outputs=model_description.outputs, @@ -87,7 +87,7 @@ def create( and weights.tensorflow_saved_model_bundle is not None ): try: - from ._tensorflow_model_adapter import TensorflowModelAdapter + from .tensorflow_backend import TensorflowModelAdapter return TensorflowModelAdapter( model_description=model_description, devices=devices @@ -96,7 +96,7 @@ def create( errors.append((wf, e)) elif wf == "onnx" and weights.onnx is not None: try: - from ._onnx_model_adapter import ONNXModelAdapter + from .onnx_backend import ONNXModelAdapter return ONNXModelAdapter( model_description=model_description, devices=devices @@ -105,7 +105,7 @@ def create( errors.append((wf, e)) elif wf == "torchscript" and weights.torchscript is not None: try: - from ._torchscript_model_adapter import TorchscriptModelAdapter + from .torchscript_backend import TorchscriptModelAdapter return TorchscriptModelAdapter( model_description=model_description, devices=devices @@ -117,13 +117,10 @@ def create( # we try to first import the keras model adapter using the separate package and, # if it is not available, try to load the one using tf try: - from ._keras import ( - KerasModelAdapter, - keras, # type: ignore - ) - - if keras is None: - from ._tensorflow_model_adapter import KerasModelAdapter + try: + from .keras_backend import KerasModelAdapter + except Exception: + from .tensorflow_backend import KerasModelAdapter return KerasModelAdapter( model_description=model_description, devices=devices @@ -134,10 +131,11 @@ def create( assert errors if len(weight_format_priority_order) == 1: assert len(errors) == 1 + wf, e = errors[0] raise ValueError( - f"The '{weight_format_priority_order[0]}' model adapter could not be created" - + f" in this environment:\n{errors[0][1].__class__.__name__}({errors[0][1]}).\n\n" - ) from errors[0][1] + f"The '{wf}' model adapter could not be created" + + f" in this environment:\n{e.__class__.__name__}({e}).\n\n" + ) from e else: error_list = "\n - ".join( @@ -165,13 +163,3 @@ def unload(self): Unload model from any devices, freeing their memory. The moder adapter should be considered unusable afterwards. """ - - -def get_weight_formats() -> List[str]: - """ - Return list of supported weight types - """ - return list(DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER) - - -create_model_adapter = ModelAdapter.create diff --git a/bioimageio/core/backend/keras.py b/bioimageio/core/backends/keras_backend.py similarity index 74% rename from bioimageio/core/backend/keras.py rename to bioimageio/core/backends/keras_backend.py index 1d273cfc..35ee79fe 100644 --- a/bioimageio/core/backend/keras.py +++ b/bioimageio/core/backends/keras_backend.py @@ -10,30 +10,22 @@ from .._settings import settings from ..digest_spec import get_axes_infos -from ..model_adapters import ModelAdapter from ..tensor import Tensor +from ._model_adapter import ModelAdapter os.environ["KERAS_BACKEND"] = settings.keras_backend # by default, we use the keras integrated with tensorflow +# TODO: check if we should prefer keras try: - import tensorflow as tf # pyright: ignore[reportMissingImports] - from tensorflow import ( # pyright: ignore[reportMissingImports] - keras, # pyright: ignore[reportUnknownVariableType] + import tensorflow as tf + from tensorflow import ( + keras, # pyright: ignore[reportUnknownVariableType,reportAttributeAccessIssue] ) - tf_version = Version(tf.__version__) # pyright: ignore[reportUnknownArgumentType] + tf_version = Version(tf.__version__) except Exception: - try: - import keras # pyright: ignore[reportMissingImports] - except Exception as e: - keras = None - keras_error = str(e) - else: - keras_error = None - tf_version = None -else: - keras_error = None + import keras class KerasModelAdapter(ModelAdapter): @@ -43,9 +35,6 @@ def __init__( model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], devices: Optional[Sequence[str]] = None, ) -> None: - if keras is None: - raise ImportError(f"failed to import keras: {keras_error}") - super().__init__() if model_description.weights.keras_hdf5 is None: raise ValueError("model has not keras_hdf5 weights specified") @@ -86,18 +75,26 @@ def __init__( def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: _result: Union[Sequence[NDArray[Any]], NDArray[Any]] - _result = self._network.predict( # pyright: ignore[reportUnknownVariableType] + _result = self._network.predict( # type: ignore *[None if t is None else t.data.data for t in input_tensors] ) if isinstance(_result, (tuple, list)): - result: Sequence[NDArray[Any]] = _result + result = _result # pyright: ignore[reportUnknownVariableType] else: result = [_result] # type: ignore - assert len(result) == len(self._output_axes) + assert len(result) == len( # pyright: ignore[reportUnknownArgumentType] + self._output_axes + ) ret: List[Optional[Tensor]] = [] ret.extend( - [Tensor(r, dims=axes) for r, axes, in zip(result, self._output_axes)] + [ + Tensor(r, dims=axes) # pyright: ignore[reportArgumentType] + for r, axes, in zip( # pyright: ignore[reportUnknownVariableType] + result, # pyright: ignore[reportUnknownArgumentType] + self._output_axes, + ) + ] ) return ret diff --git a/bioimageio/core/backends/onnx_backend.py b/bioimageio/core/backends/onnx_backend.py new file mode 100644 index 00000000..21bbcc09 --- /dev/null +++ b/bioimageio/core/backends/onnx_backend.py @@ -0,0 +1,60 @@ +import warnings +from typing import Any, List, Optional, Sequence, Union + +import onnxruntime as rt + +from bioimageio.spec._internal.type_guards import is_list, is_tuple +from bioimageio.spec.model import v0_4, v0_5 +from bioimageio.spec.utils import download + +from ..digest_spec import get_axes_infos +from ..model_adapters import ModelAdapter +from ..tensor import Tensor + + +class ONNXModelAdapter(ModelAdapter): + def __init__( + self, + *, + model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], + devices: Optional[Sequence[str]] = None, + ): + super().__init__() + self._internal_output_axes = [ + tuple(a.id for a in get_axes_infos(out)) + for out in model_description.outputs + ] + if model_description.weights.onnx is None: + raise ValueError("No ONNX weights specified for {model_description.name}") + + self._session = rt.InferenceSession( + str(download(model_description.weights.onnx.source).path) + ) + onnx_inputs = self._session.get_inputs() # type: ignore + self._input_names: List[str] = [ipt.name for ipt in onnx_inputs] # type: ignore + + 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: Optional[Tensor]) -> List[Optional[Tensor]]: + assert len(input_tensors) == len(self._input_names) + input_arrays = [None if ipt is None else ipt.data.data for ipt in input_tensors] + result: Any = self._session.run( + None, dict(zip(self._input_names, input_arrays)) + ) + if is_list(result) or is_tuple(result): + result_seq = result + else: + result_seq = [result] + + return [ + None if r is None else Tensor(r, dims=axes) + for r, axes in zip(result_seq, self._internal_output_axes) + ] + + def unload(self) -> None: + warnings.warn( + "Device management is not implemented for onnx yet, cannot unload model" + ) diff --git a/bioimageio/core/backend/pytorch.py b/bioimageio/core/backends/pytorch_backend.py similarity index 100% rename from bioimageio/core/backend/pytorch.py rename to bioimageio/core/backends/pytorch_backend.py diff --git a/bioimageio/core/backends/tensorflow_backend.py b/bioimageio/core/backends/tensorflow_backend.py new file mode 100644 index 00000000..3f9cee9d --- /dev/null +++ b/bioimageio/core/backends/tensorflow_backend.py @@ -0,0 +1,289 @@ +import zipfile +from io import TextIOWrapper +from pathlib import Path +from shutil import copyfileobj +from typing import List, Literal, Optional, Sequence, Union + +import numpy as np +import tensorflow as tf +from loguru import logger + +from bioimageio.spec.common import FileSource, ZipPath +from bioimageio.spec.model import v0_4, v0_5 +from bioimageio.spec.utils import download + +from ..digest_spec import get_axes_infos +from ..tensor import Tensor +from ._model_adapter import ModelAdapter + + +class TensorflowModelAdapterBase(ModelAdapter): + weight_format: Literal["keras_hdf5", "tensorflow_saved_model_bundle"] + + def __init__( + self, + *, + devices: Optional[Sequence[str]] = None, + weights: Union[ + v0_4.KerasHdf5WeightsDescr, + v0_4.TensorflowSavedModelBundleWeightsDescr, + v0_5.KerasHdf5WeightsDescr, + v0_5.TensorflowSavedModelBundleWeightsDescr, + ], + model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], + ): + super().__init__() + self.model_description = model_description + tf_version = v0_5.Version(tf.__version__) + model_tf_version = weights.tensorflow_version + if model_tf_version is None: + 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: + 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, + ): + logger.warning( + "The tensorflow version specified by the model does not match the installed: " + + f"{model_tf_version} != {tf_version}." + ) + + self.use_keras_api = ( + tf_version.major > 1 + or self.weight_format == KerasModelAdapter.weight_format + ) + + # TODO tf device management + if devices is not None: + logger.warning( + f"Device management is not implemented for tensorflow yet, ignoring the devices {devices}" + ) + + weight_file = self.require_unzipped(weights.source) + self._network = self._get_network(weight_file) + self._internal_output_axes = [ + tuple(a.id for a in get_axes_infos(out)) + for out in model_description.outputs + ] + + # TODO: check how to load tf weights without unzipping + def require_unzipped(self, weight_file: FileSource): + local_weights_file = download(weight_file).path + if isinstance(local_weights_file, ZipPath): + # weights file is in a bioimageio zip package + out_path = ( + Path("bioimageio_unzipped_tf_weights") / local_weights_file.filename + ) + with local_weights_file.open("rb") as src, out_path.open("wb") as dst: + assert not isinstance(src, TextIOWrapper) + copyfileobj(src, dst) + + local_weights_file = out_path + + if zipfile.is_zipfile(local_weights_file): + # weights file itself is a zipfile + out_path = local_weights_file.with_suffix(".unzipped") + with zipfile.ZipFile(local_weights_file, "r") as f: + f.extractall(out_path) + + return out_path + else: + return local_weights_file + + def _get_network( # pyright: ignore[reportUnknownParameterType] + self, weight_file: FileSource + ): + weight_file = self.require_unzipped(weight_file) + assert tf is not None + if self.use_keras_api: + try: + return tf.keras.layers.TFSMLayer( # pyright: ignore[reportAttributeAccessIssue,reportUnknownVariableType] + weight_file, + call_endpoint="serve", + ) + except Exception as e: + try: + return tf.keras.layers.TFSMLayer( # pyright: ignore[reportAttributeAccessIssue,reportUnknownVariableType] + weight_file, call_endpoint="serving_default" + ) + 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) + + # 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( # pyright: ignore[reportUnknownParameterType] + self, *input_tensors: Optional[Tensor] + ): + assert tf is not None + input_keys = [ + ipt.name if isinstance(ipt, v0_4.InputTensorDescr) else ipt.id + for ipt in self.model_description.inputs + ] + output_keys = [ + out.name if isinstance(out, v0_4.OutputTensorDescr) else out.id + for out in self.model_description.outputs + ] + # TODO read from spec + tag = ( # pyright: ignore[reportUnknownVariableType] + tf.saved_model.tag_constants.SERVING # pyright: ignore[reportAttributeAccessIssue] + ) + signature_key = ( # pyright: ignore[reportUnknownVariableType] + tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY # pyright: ignore[reportAttributeAccessIssue] + ) + + graph = tf.Graph() + with graph.as_default(): + with tf.Session( # pyright: ignore[reportAttributeAccessIssue] + graph=graph + ) as sess: # pyright: ignore[reportUnknownVariableType] + # load the model and the signature + graph_def = tf.saved_model.loader.load( # pyright: ignore[reportUnknownVariableType,reportAttributeAccessIssue] + sess, [tag], self._network + ) + signature = ( # pyright: ignore[reportUnknownVariableType] + graph_def.signature_def + ) + + # get the tensors into the graph + in_names = [ # pyright: ignore[reportUnknownVariableType] + signature[signature_key].inputs[key].name for key in input_keys + ] + out_names = [ # pyright: ignore[reportUnknownVariableType] + signature[signature_key].outputs[key].name for key in output_keys + ] + in_tensors = [ + graph.get_tensor_by_name( + name # pyright: ignore[reportUnknownArgumentType] + ) + for name in in_names # pyright: ignore[reportUnknownVariableType] + ] + out_tensors = [ + graph.get_tensor_by_name( + name # pyright: ignore[reportUnknownArgumentType] + ) + for name in out_names # pyright: ignore[reportUnknownVariableType] + ] + + # run prediction + res = sess.run( # pyright: ignore[reportUnknownVariableType] + dict( + zip( + out_names, # pyright: ignore[reportUnknownArgumentType] + out_tensors, + ) + ), + dict( + zip( + in_tensors, + [None if t is None else t.data for t in input_tensors], + ) + ), + ) + # from dict to list of tensors + res = [ # pyright: ignore[reportUnknownVariableType] + res[out] + for out in out_names # pyright: ignore[reportUnknownVariableType] + ] + + return res # pyright: ignore[reportUnknownVariableType] + + def _forward_keras( # pyright: ignore[reportUnknownParameterType] + self, *input_tensors: Optional[Tensor] + ): + assert self.use_keras_api + assert not isinstance(self._network, str) + assert tf is not None + tf_tensor = [ + None if ipt is None else tf.convert_to_tensor(ipt) for ipt in input_tensors + ] + + result = self._network(*tf_tensor) # pyright: ignore[reportUnknownVariableType] + + assert isinstance(result, dict) + + # TODO: Use RDF's `outputs[i].id` here + result = list( # pyright: ignore[reportUnknownVariableType] + result.values() # pyright: ignore[reportUnknownArgumentType] + ) + + return [ # pyright: ignore[reportUnknownVariableType] + (None if r is None else r if isinstance(r, np.ndarray) else r.numpy()) + for r in result # pyright: ignore[reportUnknownVariableType] + ] + + def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: + if self.use_keras_api: + result = self._forward_keras( # pyright: ignore[reportUnknownVariableType] + *input_tensors + ) + else: + result = self._forward_tf( # pyright: ignore[reportUnknownVariableType] + *input_tensors + ) + + return [ + ( + None + if r is None + else Tensor(r, dims=axes) # pyright: ignore[reportUnknownArgumentType] + ) + for r, axes in zip( # pyright: ignore[reportUnknownVariableType] + result, # pyright: ignore[reportUnknownArgumentType] + self._internal_output_axes, + ) + ] + + def unload(self) -> None: + logger.warning( + "Device management is not implemented for keras yet, cannot unload model" + ) + + +class TensorflowModelAdapter(TensorflowModelAdapterBase): + weight_format = "tensorflow_saved_model_bundle" + + def __init__( + self, + *, + model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], + devices: Optional[Sequence[str]] = None, + ): + if model_description.weights.tensorflow_saved_model_bundle is None: + raise ValueError("missing tensorflow_saved_model_bundle weights") + + super().__init__( + devices=devices, + weights=model_description.weights.tensorflow_saved_model_bundle, + model_description=model_description, + ) + + +class KerasModelAdapter(TensorflowModelAdapterBase): + weight_format = "keras_hdf5" + + def __init__( + self, + *, + model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], + devices: Optional[Sequence[str]] = None, + ): + if model_description.weights.keras_hdf5 is None: + raise ValueError("missing keras_hdf5 weights") + + super().__init__( + model_description=model_description, + devices=devices, + weights=model_description.weights.keras_hdf5, + ) diff --git a/bioimageio/core/backends/torchscript_backend.py b/bioimageio/core/backends/torchscript_backend.py new file mode 100644 index 00000000..d1882180 --- /dev/null +++ b/bioimageio/core/backends/torchscript_backend.py @@ -0,0 +1,79 @@ +import gc +import warnings +from typing import Any, List, Optional, Sequence, Union + +import torch + +from bioimageio.spec._internal.type_guards import is_list, is_ndarray, is_tuple +from bioimageio.spec.model import v0_4, v0_5 +from bioimageio.spec.utils import download + +from ..digest_spec import get_axes_infos +from ..model_adapters import ModelAdapter +from ..tensor import Tensor + + +class TorchscriptModelAdapter(ModelAdapter): + def __init__( + self, + *, + model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], + devices: Optional[Sequence[str]] = None, + ): + super().__init__() + if model_description.weights.torchscript is None: + raise ValueError( + f"No torchscript weights found for model {model_description.name}" + ) + + weight_path = download(model_description.weights.torchscript.source).path + 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._model = self._model.eval() + self._internal_output_axes = [ + tuple(a.id for a in get_axes_infos(out)) + for out in model_description.outputs + ] + + def forward(self, *batch: Optional[Tensor]) -> List[Optional[Tensor]]: + with torch.no_grad(): + torch_tensor = [ + None if b is None else torch.from_numpy(b.data.data).to(self.devices[0]) + for b in batch + ] + _result: Any = self._model.forward(*torch_tensor) + if is_list(_result) or is_tuple(_result): + result: Sequence[Any] = _result + else: + result = [_result] + + result = [ + ( + None + if r is None + else r.cpu().numpy() if isinstance(r, torch.Tensor) else r + ) + for r in result + ] + + assert len(result) == len(self._internal_output_axes) + return [ + None if r is None else Tensor(r, dims=axes) if is_ndarray(r) else r + 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 diff --git a/bioimageio/core/model_adapters.py b/bioimageio/core/model_adapters.py index 86fcfe4b..db92d013 100644 --- a/bioimageio/core/model_adapters.py +++ b/bioimageio/core/model_adapters.py @@ -1,8 +1,22 @@ -from ._create_model_adapter import create_model_adapter -from ._model_adapter import ModelAdapter, get_weight_formats +"""DEPRECATED""" + +from typing import List + +from .backends._model_adapter import ( + DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER, + ModelAdapter, + create_model_adapter, +) __all__ = [ "ModelAdapter", "create_model_adapter", "get_weight_formats", ] + + +def get_weight_formats() -> List[str]: + """ + Return list of supported weight types + """ + return list(DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER) diff --git a/bioimageio/core/model_adapters/_pytorch_model_adapter.py b/bioimageio/core/model_adapters/_pytorch_model_adapter.py deleted file mode 100644 index e69de29b..00000000 From f52a894231fc81792045033d726c6eda0db6e82b Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 6 Dec 2024 16:42:16 +0100 Subject: [PATCH 24/47] sort tests --- tests/test_proc_ops.py | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/tests/test_proc_ops.py b/tests/test_proc_ops.py index e408d220..0b93f08b 100644 --- a/tests/test_proc_ops.py +++ b/tests/test_proc_ops.py @@ -105,6 +105,22 @@ def test_zero_mean_unit_variance_fixed(tid: MemberId): xr.testing.assert_allclose(expected, sample.members[tid].data, rtol=1e-5, atol=1e-7) +def test_zero_mean_unit_variance_fixed2(tid: MemberId): + from bioimageio.core.proc_ops import FixedZeroMeanUnitVariance + + np_data = np.arange(9).reshape(3, 3) + mean = float(np_data.mean()) + std = float(np_data.mean()) + eps = 1.0e-7 + op = FixedZeroMeanUnitVariance(tid, tid, mean=mean, std=std, eps=eps) + + data = xr.DataArray(np_data, dims=("x", "y")) + sample = Sample(members={tid: Tensor.from_xarray(data)}, stat={}, id=None) + expected = xr.DataArray((np_data - mean) / (std + eps), dims=("x", "y")) + op(sample) + xr.testing.assert_allclose(expected, sample.members[tid].data, rtol=1e-5, atol=1e-7) + + def test_zero_mean_unit_across_axes(tid: MemberId): from bioimageio.core.proc_ops import ZeroMeanUnitVariance @@ -126,22 +142,6 @@ def test_zero_mean_unit_across_axes(tid: MemberId): xr.testing.assert_allclose(expected, sample.members[tid].data, rtol=1e-5, atol=1e-7) -def test_zero_mean_unit_variance_fixed2(tid: MemberId): - from bioimageio.core.proc_ops import FixedZeroMeanUnitVariance - - np_data = np.arange(9).reshape(3, 3) - mean = float(np_data.mean()) - std = float(np_data.mean()) - eps = 1.0e-7 - op = FixedZeroMeanUnitVariance(tid, tid, mean=mean, std=std, eps=eps) - - data = xr.DataArray(np_data, dims=("x", "y")) - sample = Sample(members={tid: Tensor.from_xarray(data)}, stat={}, id=None) - expected = xr.DataArray((np_data - mean) / (std + eps), dims=("x", "y")) - op(sample) - xr.testing.assert_allclose(expected, sample.members[tid].data, rtol=1e-5, atol=1e-7) - - def test_binarize(tid: MemberId): from bioimageio.core.proc_ops import Binarize From 523c54b4f60a3076df7ddd24ed4ac310a6b8a875 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 9 Dec 2024 13:43:39 +0100 Subject: [PATCH 25/47] add create_model_adapter --- bioimageio/core/backends/_model_adapter.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/bioimageio/core/backends/_model_adapter.py b/bioimageio/core/backends/_model_adapter.py index 66153f09..99919aac 100644 --- a/bioimageio/core/backends/_model_adapter.py +++ b/bioimageio/core/backends/_model_adapter.py @@ -163,3 +163,6 @@ def unload(self): Unload model from any devices, freeing their memory. The moder adapter should be considered unusable afterwards. """ + + +create_model_adapter = ModelAdapter.create From d438a123597cc707124a3b6e1569d317fec1b3ef Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 9 Dec 2024 13:43:52 +0100 Subject: [PATCH 26/47] pin pyright --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index af913c1d..95414371 100644 --- a/setup.py +++ b/setup.py @@ -61,7 +61,7 @@ "pre-commit", "pdoc", "psutil", # parallel pytest with 'pytest -n auto' - "pyright", + "pyright==1.1.390", "pytest-cov", "pytest-xdist", # parallel pytest "pytest", From f9a1a67a7e84e17f6390350d712ae7075d70929d Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 9 Dec 2024 14:43:02 +0100 Subject: [PATCH 27/47] continue refactor of weight converters and backends --- bioimageio/core/_resource_tests.py | 12 +- bioimageio/core/backends/pytorch_backend.py | 159 +++--- .../core/backends/tensorflow_backend.py | 40 +- bioimageio/core/io.py | 30 +- bioimageio/core/test_bioimageio_collection.py | 60 --- bioimageio/core/weight_converters.py | 492 ------------------ .../weight_converters/keras_to_tensorflow.py | 184 +++++++ .../core/weight_converters/pytorch_to_onnx.py | 17 +- .../pytorch_to_torchscript.py | 154 ++++++ tests/conftest.py | 4 +- tests/test_weight_converters.py | 19 +- tests/utils.py | 2 + 12 files changed, 476 insertions(+), 697 deletions(-) delete mode 100644 bioimageio/core/test_bioimageio_collection.py delete mode 100644 bioimageio/core/weight_converters.py create mode 100644 bioimageio/core/weight_converters/keras_to_tensorflow.py create mode 100644 bioimageio/core/weight_converters/pytorch_to_torchscript.py diff --git a/bioimageio/core/_resource_tests.py b/bioimageio/core/_resource_tests.py index e6675b73..0dae30ff 100644 --- a/bioimageio/core/_resource_tests.py +++ b/bioimageio/core/_resource_tests.py @@ -428,7 +428,7 @@ def _test_model_inference( rtol: float, ) -> None: test_name = f"Reproduce test outputs from test inputs ({weight_format})" - logger.info("starting '{}'", test_name) + logger.debug("starting '{}'", test_name) error: Optional[str] = None tb: List[str] = [] @@ -516,11 +516,13 @@ def _test_model_inference_parametrized( # 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)) + test_cases: Set[Tuple[BatchSize, v0_5.ParameterizedSize_N]] = { + (b, n) for b, n in product(sorted(batch_sizes), sorted(ns)) } logger.info( - "Testing inference with {} different input tensor sizes", len(test_cases) + "Testing inference with {} different inputs (B, N): {}", + len(test_cases), + test_cases, ) def generate_test_cases(): @@ -534,7 +536,7 @@ def get_ns(n: int): if isinstance(a.size, v0_5.ParameterizedSize) } - for n, batch_size in sorted(test_cases): + for batch_size, n in sorted(test_cases): input_target_sizes, expected_output_sizes = model.get_axis_sizes( get_ns(n), batch_size=batch_size ) diff --git a/bioimageio/core/backends/pytorch_backend.py b/bioimageio/core/backends/pytorch_backend.py index 1992f406..74e59f30 100644 --- a/bioimageio/core/backends/pytorch_backend.py +++ b/bioimageio/core/backends/pytorch_backend.py @@ -34,12 +34,12 @@ def __init__( ): super().__init__() self.output_dims = [tuple(a.id for a in get_axes_infos(out)) for out in outputs] - devices = self.get_devices(devices) - self._network = self.get_network(weights, load_state=True, devices=devices) + devices = get_devices(devices) + self._model = load_torch_model(weights, load_state=True, devices=devices) if mode == "eval": - self._network = self._network.eval() + self._model = self._model.eval() elif mode == "train": - self._network = self._network.train() + self._model = self._model.train() else: assert_never(mode) @@ -63,7 +63,7 @@ def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: (None if t is None else t.to(self._primary_device)) for t in tensors ] result: Union[Tuple[Any, ...], List[Any], Any] - result = self._network(*tensors) + result = self._model(*tensors) if not isinstance(result, (tuple, list)): result = [result] @@ -86,91 +86,84 @@ def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: ] def unload(self) -> None: - del self._network + del self._model _ = gc.collect() # deallocate memory assert torch is not None torch.cuda.empty_cache() # release reserved memory - @classmethod - def get_network( - cls, - weight_spec: Union[ - v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr - ], - *, - load_state: bool = False, - devices: Optional[Sequence[Union[str, torch.device]]] = None, - ) -> nn.Module: - arch = import_callable( - weight_spec.architecture, - sha256=( - weight_spec.architecture_sha256 - if isinstance(weight_spec, v0_4.PytorchStateDictWeightsDescr) - else weight_spec.sha256 - ), - ) - model_kwargs = ( - weight_spec.kwargs + +def load_torch_model( + weight_spec: Union[ + v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr + ], + *, + load_state: bool = False, + devices: Optional[Sequence[Union[str, torch.device]]] = None, +) -> nn.Module: + arch = import_callable( + weight_spec.architecture, + sha256=( + weight_spec.architecture_sha256 if isinstance(weight_spec, v0_4.PytorchStateDictWeightsDescr) - else weight_spec.architecture.kwargs + else weight_spec.sha256 + ), + ) + model_kwargs = ( + weight_spec.kwargs + if isinstance(weight_spec, v0_4.PytorchStateDictWeightsDescr) + else weight_spec.architecture.kwargs + ) + network = arch(**model_kwargs) + if not isinstance(network, nn.Module): + raise ValueError( + f"calling {weight_spec.architecture.callable} did not return a torch.nn.Module" ) - network = arch(**model_kwargs) - if not isinstance(network, nn.Module): - raise ValueError( - f"calling {weight_spec.architecture.callable} did not return a torch.nn.Module" - ) - if load_state or devices: - use_devices = cls.get_devices(devices) - network = network.to(use_devices[0]) - if load_state: - network = cls.load_state( - network, - path=download(weight_spec).path, - devices=use_devices, - ) - return network - - @staticmethod - def load_state( - network: nn.Module, - path: Union[Path, ZipPath], - devices: Sequence[torch.device], - ) -> nn.Module: - network = network.to(devices[0]) - with path.open("rb") as f: - assert not isinstance(f, TextIOWrapper) - state = torch.load(f, map_location=devices[0]) - - incompatible = network.load_state_dict(state) - if incompatible.missing_keys: - logger.warning("Missing state dict keys: {}", incompatible.missing_keys) - - if incompatible.unexpected_keys: - logger.warning( - "Unexpected state dict keys: {}", incompatible.unexpected_keys + if load_state or devices: + use_devices = get_devices(devices) + network = network.to(use_devices[0]) + if load_state: + network = load_torch_state_dict( + network, + path=download(weight_spec).path, + devices=use_devices, ) - return network - - @staticmethod - def get_devices( - devices: Optional[Sequence[Union[torch.device, str]]] = None, - ) -> List[torch.device]: - if not devices: - torch_devices = [ - ( - torch.device("cuda") - if torch.cuda.is_available() - else torch.device("cpu") - ) - ] - else: - torch_devices = [torch.device(d) for d in devices] + return network + + +def load_torch_state_dict( + model: nn.Module, + path: Union[Path, ZipPath], + devices: Sequence[torch.device], +) -> nn.Module: + model = model.to(devices[0]) + with path.open("rb") as f: + assert not isinstance(f, TextIOWrapper) + state = torch.load(f, map_location=devices[0]) + + incompatible = model.load_state_dict(state) + if incompatible.missing_keys: + logger.warning("Missing state dict keys: {}", incompatible.missing_keys) + + if incompatible.unexpected_keys: + logger.warning("Unexpected state dict keys: {}", incompatible.unexpected_keys) + return model + + +def get_devices( + devices: Optional[Sequence[Union[torch.device, str]]] = None, +) -> List[torch.device]: + if not devices: + torch_devices = [ + (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")) + ] + else: + torch_devices = [torch.device(d) for d in devices] - if len(torch_devices) > 1: - warnings.warn( - f"Multiple devices for single pytorch model not yet implemented; ignoring {torch_devices[1:]}" - ) - torch_devices = torch_devices[:1] + if len(torch_devices) > 1: + warnings.warn( + f"Multiple devices for single pytorch model not yet implemented; ignoring {torch_devices[1:]}" + ) + torch_devices = torch_devices[:1] - return torch_devices + return torch_devices diff --git a/bioimageio/core/backends/tensorflow_backend.py b/bioimageio/core/backends/tensorflow_backend.py index 3f9cee9d..94a8165f 100644 --- a/bioimageio/core/backends/tensorflow_backend.py +++ b/bioimageio/core/backends/tensorflow_backend.py @@ -1,16 +1,13 @@ -import zipfile -from io import TextIOWrapper from pathlib import Path -from shutil import copyfileobj from typing import List, Literal, Optional, Sequence, Union import numpy as np import tensorflow as tf from loguru import logger -from bioimageio.spec.common import FileSource, ZipPath +from bioimageio.core.io import ensure_unzipped +from bioimageio.spec.common import FileSource from bioimageio.spec.model import v0_4, v0_5 -from bioimageio.spec.utils import download from ..digest_spec import get_axes_infos from ..tensor import Tensor @@ -65,41 +62,22 @@ def __init__( f"Device management is not implemented for tensorflow yet, ignoring the devices {devices}" ) - weight_file = self.require_unzipped(weights.source) + # TODO: check how to load tf weights without unzipping + weight_file = ensure_unzipped( + weights.source, Path("bioimageio_unzipped_tf_weights") + ) self._network = self._get_network(weight_file) self._internal_output_axes = [ tuple(a.id for a in get_axes_infos(out)) for out in model_description.outputs ] - # TODO: check how to load tf weights without unzipping - def require_unzipped(self, weight_file: FileSource): - local_weights_file = download(weight_file).path - if isinstance(local_weights_file, ZipPath): - # weights file is in a bioimageio zip package - out_path = ( - Path("bioimageio_unzipped_tf_weights") / local_weights_file.filename - ) - with local_weights_file.open("rb") as src, out_path.open("wb") as dst: - assert not isinstance(src, TextIOWrapper) - copyfileobj(src, dst) - - local_weights_file = out_path - - if zipfile.is_zipfile(local_weights_file): - # weights file itself is a zipfile - out_path = local_weights_file.with_suffix(".unzipped") - with zipfile.ZipFile(local_weights_file, "r") as f: - f.extractall(out_path) - - return out_path - else: - return local_weights_file - def _get_network( # pyright: ignore[reportUnknownParameterType] self, weight_file: FileSource ): - weight_file = self.require_unzipped(weight_file) + weight_file = ensure_unzipped( + weight_file, Path("bioimageio_unzipped_tf_weights") + ) assert tf is not None if self.use_keras_api: try: diff --git a/bioimageio/core/io.py b/bioimageio/core/io.py index ee60a67a..001db539 100644 --- a/bioimageio/core/io.py +++ b/bioimageio/core/io.py @@ -1,6 +1,9 @@ import collections.abc import warnings +import zipfile +from io import TextIOWrapper from pathlib import Path, PurePosixPath +from shutil import copyfileobj from typing import Any, Mapping, Optional, Sequence, Tuple, Union import h5py @@ -10,7 +13,8 @@ from numpy.typing import NDArray from pydantic import BaseModel, ConfigDict, TypeAdapter -from bioimageio.spec.utils import load_array, save_array +from bioimageio.spec.common import FileSource, ZipPath +from bioimageio.spec.utils import download, load_array, save_array from .axis import AxisLike from .common import PerMember @@ -176,3 +180,27 @@ def save_dataset_stat(stat: Mapping[DatasetMeasure, MeasureValue], path: Path): def load_dataset_stat(path: Path): seq = _stat_adapter.validate_json(path.read_bytes()) return {e.measure: e.value for e in seq} + + +def ensure_unzipped(source: Union[FileSource, ZipPath], folder: Path): + """unzip a (downloaded) **source** to a file in **folder** if source is a zip archive. + Always returns the path to the unzipped source (maybe source itself)""" + local_weights_file = download(source).path + if isinstance(local_weights_file, ZipPath): + # source is inside a zip archive + out_path = folder / local_weights_file.filename + with local_weights_file.open("rb") as src, out_path.open("wb") as dst: + assert not isinstance(src, TextIOWrapper) + copyfileobj(src, dst) + + local_weights_file = out_path + + if zipfile.is_zipfile(local_weights_file): + # source itself is a zipfile + out_path = folder / local_weights_file.with_suffix(".unzipped").name + with zipfile.ZipFile(local_weights_file, "r") as f: + f.extractall(out_path) + + return out_path + else: + return local_weights_file diff --git a/bioimageio/core/test_bioimageio_collection.py b/bioimageio/core/test_bioimageio_collection.py deleted file mode 100644 index 2cf9ced0..00000000 --- a/bioimageio/core/test_bioimageio_collection.py +++ /dev/null @@ -1,60 +0,0 @@ -from typing import Any, Collection, Dict, Iterable, Mapping, Tuple - -import pytest -import requests -from pydantic import HttpUrl - -from bioimageio.spec import InvalidDescr -from bioimageio.spec.common import Sha256 -from tests.utils import ParameterSet, expensive_test - -BASE_URL = "https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/" - - -def _get_latest_rdf_sources(): - entries: Any = requests.get(BASE_URL + "all_versions.json").json()["entries"] - ret: Dict[str, Tuple[HttpUrl, Sha256]] = {} - for entry in entries: - version = entry["versions"][0] - ret[f"{entry['concept']}/{version['v']}"] = ( - HttpUrl(version["source"]), # pyright: ignore[reportCallIssue] - Sha256(version["sha256"]), - ) - - return ret - - -ALL_LATEST_RDF_SOURCES: Mapping[str, Tuple[HttpUrl, Sha256]] = _get_latest_rdf_sources() - - -def yield_bioimageio_yaml_urls() -> Iterable[ParameterSet]: - for descr_url, sha in ALL_LATEST_RDF_SOURCES.values(): - key = ( - str(descr_url) - .replace(BASE_URL, "") - .replace("/files/rdf.yaml", "") - .replace("/files/bioimageio.yaml", "") - ) - yield pytest.param(descr_url, sha, key, id=key) - - -KNOWN_INVALID: Collection[str] = set() - - -@expensive_test -@pytest.mark.parametrize("descr_url,sha,key", list(yield_bioimageio_yaml_urls())) -def test_rdf( - descr_url: HttpUrl, - sha: Sha256, - key: str, -): - if key in KNOWN_INVALID: - pytest.skip("known failure") - - from bioimageio.core import load_description_and_test - - descr = load_description_and_test(descr_url, sha256=sha) - assert not isinstance(descr, InvalidDescr) - assert ( - descr.validation_summary.status == "passed" - ), descr.validation_summary.format() diff --git a/bioimageio/core/weight_converters.py b/bioimageio/core/weight_converters.py deleted file mode 100644 index 6e0d06ec..00000000 --- a/bioimageio/core/weight_converters.py +++ /dev/null @@ -1,492 +0,0 @@ -# type: ignore # TODO: type -from __future__ import annotations - -import abc -from bioimageio.spec.model.v0_5 import WeightsEntryDescrBase -from typing import Any, List, Sequence, cast, Union -from typing_extensions import assert_never -import numpy as np -from numpy.testing import assert_array_almost_equal -from bioimageio.spec.model import v0_4, v0_5 -from torch.jit import ScriptModule -from bioimageio.core.digest_spec import get_test_inputs, get_member_id -from bioimageio.core.model_adapters._pytorch_model_adapter import PytorchModelAdapter -import os -import shutil -from pathlib import Path -from typing import no_type_check -from zipfile import ZipFile -from bioimageio.spec._internal.version_type import Version -from bioimageio.spec._internal.io_utils import download - -try: - import torch -except ImportError: - torch = None - -try: - import tensorflow.saved_model -except Exception: - tensorflow = None - - -# additional convenience for pytorch state dict, eventually we want this in python-bioimageio too -# and for each weight format -def load_torch_model( # pyright: ignore[reportUnknownParameterType] - node: Union[v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr], -): - assert torch is not None - model = ( # pyright: ignore[reportUnknownVariableType] - PytorchModelAdapter.get_network(node) - ) - state = torch.load(download(node.source).path, map_location="cpu") - model.load_state_dict(state) # FIXME: check incompatible keys? - return model.eval() # pyright: ignore[reportUnknownVariableType] - - -class WeightConverter(abc.ABC): - @abc.abstractmethod - def convert( - self, model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], output_path: Path - ) -> WeightsEntryDescrBase: - raise NotImplementedError - - -class Pytorch2Onnx(WeightConverter): - def __init__(self): - super().__init__() - assert torch is not None - - def convert( - self, - model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], - output_path: Path, - use_tracing: bool = True, - test_decimal: int = 4, - verbose: bool = False, - opset_version: int = 15, - ) -> v0_5.OnnxWeightsDescr: - """ - Convert model weights from the PyTorch state_dict format to the ONNX format. - - Args: - model_descr (Union[v0_4.ModelDescr, v0_5.ModelDescr]): - The model description object that contains the model and its weights. - output_path (Path): - The file path where the ONNX model will be saved. - use_tracing (bool, optional): - Whether to use tracing or scripting to export the ONNX format. Defaults to True. - test_decimal (int, optional): - The decimal precision for comparing the results between the original and converted models. - This is used in the `assert_array_almost_equal` function to check if the outputs match. - Defaults to 4. - verbose (bool, optional): - If True, will print out detailed information during the ONNX export process. Defaults to False. - opset_version (int, optional): - The ONNX opset version to use for the export. Defaults to 15. - - Raises: - ValueError: - If the provided model does not have weights in the PyTorch state_dict format. - ImportError: - If ONNX Runtime is not available for checking the exported ONNX model. - ValueError: - If the results before and after weights conversion do not agree. - - Returns: - v0_5.OnnxWeightsDescr: - A descriptor object that contains information about the exported ONNX weights. - """ - - state_dict_weights_descr = model_descr.weights.pytorch_state_dict - if state_dict_weights_descr is None: - raise ValueError( - "The provided model does not have weights in the pytorch state dict format" - ) - - assert torch is not None - with torch.no_grad(): - sample = get_test_inputs(model_descr) - input_data = [ - sample.members[get_member_id(ipt)].data.data - for ipt in model_descr.inputs - ] - input_tensors = [torch.from_numpy(ipt) for ipt in input_data] - model = load_torch_model(state_dict_weights_descr) - - expected_tensors = model(*input_tensors) - if isinstance(expected_tensors, torch.Tensor): - expected_tensors = [expected_tensors] - expected_outputs: List[np.ndarray[Any, Any]] = [ - out.numpy() for out in expected_tensors - ] - - if use_tracing: - torch.onnx.export( - model, - ( - tuple(input_tensors) - if len(input_tensors) > 1 - else input_tensors[0] - ), - str(output_path), - verbose=verbose, - opset_version=opset_version, - ) - else: - raise NotImplementedError - - try: - import onnxruntime as rt # pyright: ignore [reportMissingTypeStubs] - except ImportError: - 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)) - onnx_input_node_args = cast( - List[Any], sess.get_inputs() - ) # fixme: remove cast, try using rt.NodeArg instead of Any - onnx_inputs = { - input_name.name: inp - for input_name, inp in zip(onnx_input_node_args, input_data) - } - outputs = cast( - Sequence[np.ndarray[Any, Any]], sess.run(None, onnx_inputs) - ) # FIXME: remove cast - - try: - for exp, out in zip(expected_outputs, outputs): - assert_array_almost_equal(exp, out, decimal=test_decimal) - except AssertionError as e: - 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 - ) - - -class Pytorch2Torchscipt(WeightConverter): - def __init__(self): - super().__init__() - assert torch is not None - - def convert( - self, - model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], - output_path: Path, - use_tracing: bool = True, - ) -> v0_5.TorchscriptWeightsDescr: - """ - Convert model weights from the PyTorch `state_dict` format to TorchScript. - - Args: - model_descr (Union[v0_4.ModelDescr, v0_5.ModelDescr]): - The model description object that contains the model and its weights in the PyTorch `state_dict` format. - output_path (Path): - The file path where the TorchScript model will be saved. - use_tracing (bool): - Whether to use tracing or scripting to export the TorchScript format. - - `True`: Use tracing, which is recommended for models with straightforward control flow. - - `False`: Use scripting, which is better for models with dynamic control flow (e.g., loops, conditionals). - - Raises: - ValueError: - If the provided model does not have weights in the PyTorch `state_dict` format. - - Returns: - v0_5.TorchscriptWeightsDescr: - A descriptor object that contains information about the exported TorchScript weights. - """ - state_dict_weights_descr = model_descr.weights.pytorch_state_dict - if state_dict_weights_descr is None: - raise ValueError( - "The provided model does not have weights in the pytorch state dict format" - ) - - input_data = model_descr.get_input_test_arrays() - - 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) - scripted_module: ScriptModule = ( - torch.jit.trace(model, input_data) - if use_tracing - else torch.jit.script(model) - ) - self._check_predictions( - model=model, - scripted_model=scripted_module, - model_spec=model_descr, - input_data=input_data, - ) - - scripted_module.save(str(output_path)) - - return v0_5.TorchscriptWeightsDescr( - source=output_path, - pytorch_version=Version(torch.__version__), - parent="pytorch_state_dict", - ) - - def _check_predictions( - self, - model: Any, - scripted_model: Any, - model_spec: v0_4.ModelDescr | v0_5.ModelDescr, - input_data: Sequence[torch.Tensor], - ): - assert torch is not None - - def _check(input_: Sequence[torch.Tensor]) -> None: - expected_tensors = model(*input_) - if isinstance(expected_tensors, torch.Tensor): - expected_tensors = [expected_tensors] - expected_outputs: List[np.ndarray[Any, Any]] = [ - out.numpy() for out in expected_tensors - ] - - output_tensors = scripted_model(*input_) - if isinstance(output_tensors, torch.Tensor): - output_tensors = [output_tensors] - outputs: List[np.ndarray[Any, Any]] = [ - out.numpy() for out in output_tensors - ] - - try: - for exp, out in zip(expected_outputs, outputs): - assert_array_almost_equal(exp, out, decimal=4) - except AssertionError as e: - raise ValueError( - f"Results before and after weights conversion do not agree:\n {str(e)}" - ) - - _check(input_data) - - if len(model_spec.inputs) > 1: - return # FIXME: why don't we check multiple inputs? - - input_descr = model_spec.inputs[0] - if isinstance(input_descr, v0_4.InputTensorDescr): - if not isinstance(input_descr.shape, v0_4.ParameterizedInputShape): - return - min_shape = input_descr.shape.min - step = input_descr.shape.step - else: - min_shape: List[int] = [] - step: List[int] = [] - for axis in input_descr.axes: - if isinstance(axis.size, v0_5.ParameterizedSize): - min_shape.append(axis.size.min) - step.append(axis.size.step) - elif isinstance(axis.size, int): - min_shape.append(axis.size) - step.append(0) - elif axis.size is None: - raise NotImplementedError( - f"Can't verify inputs that don't specify their shape fully: {axis}" - ) - elif isinstance(axis.size, v0_5.SizeReference): - raise NotImplementedError(f"Can't handle axes like '{axis}' yet") - else: - assert_never(axis.size) - - 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 step_dim == 0 - else slice(0, max_dim - step_factor * step_dim, 1) - ) - 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]) - - -class Tensorflow2Bundled(WeightConverter): - def __init__(self): - super().__init__() - assert tensorflow is not None - - def convert( - self, model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], output_path: Path - ) -> v0_5.TensorflowSavedModelBundleWeightsDescr: - """ - Convert model weights from the 'keras_hdf5' format to the 'tensorflow_saved_model_bundle' format. - - This method handles the conversion of Keras HDF5 model weights into a TensorFlow SavedModel bundle, - which is the recommended format for deploying TensorFlow models. The method supports both TensorFlow 1.x - and 2.x versions, with appropriate checks to ensure compatibility. - - Adapted from: - https://github.com/deepimagej/pydeepimagej/blob/5aaf0e71f9b04df591d5ca596f0af633a7e024f5/pydeepimagej/yaml/create_config.py - - Args: - model_descr (Union[v0_4.ModelDescr, v0_5.ModelDescr]): - The bioimage.io model description containing the model's metadata and weights. - output_path (Path): - The directory where the TensorFlow SavedModel bundle will be saved. - This path must not already exist and, if necessary, will be zipped into a .zip file. - use_tracing (bool): - Placeholder argument; currently not used in this method but required to match the abstract method signature. - - Raises: - ValueError: - - If the specified `output_path` already exists. - - If the Keras HDF5 weights are missing in the model description. - RuntimeError: - If there is a mismatch between the TensorFlow version used by the model and the version installed. - NotImplementedError: - If the model has multiple inputs or outputs and TensorFlow 1.x is being used. - - Returns: - v0_5.TensorflowSavedModelBundleWeightsDescr: - A descriptor object containing information about the converted TensorFlow SavedModel bundle. - """ - assert tensorflow is not None - tf_major_ver = int(tensorflow.__version__.split(".")[0]) - - if output_path.suffix == ".zip": - output_path = output_path.with_suffix("") - zip_weights = True - else: - zip_weights = False - - if output_path.exists(): - raise ValueError(f"The ouptut directory at {output_path} must not exist.") - - if model_descr.weights.keras_hdf5 is None: - raise ValueError("Missing Keras Hdf5 weights to convert from.") - - weight_spec = model_descr.weights.keras_hdf5 - weight_path = download(weight_spec.source).path - - if weight_spec.tensorflow_version: - model_tf_major_ver = int(weight_spec.tensorflow_version.major) - if model_tf_major_ver != tf_major_ver: - raise RuntimeError( - f"Tensorflow major versions of model {model_tf_major_ver} is not {tf_major_ver}" - ) - - if tf_major_ver == 1: - if len(model_descr.inputs) != 1 or len(model_descr.outputs) != 1: - raise NotImplementedError( - "Weight conversion for models with multiple inputs or outputs is not yet implemented." - ) - return self._convert_tf1( - weight_path, - output_path, - model_descr.inputs[0].id, - model_descr.outputs[0].id, - zip_weights, - ) - else: - return self._convert_tf2(weight_path, output_path, zip_weights) - - def _convert_tf2( - self, 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 tensorflow import keras - except Exception: - # if the above fails try to export with the standalone keras - import keras - - model = keras.models.load_model(keras_weight_path) - keras.models.save_model(model, output_path) - - if zip_weights: - output_path = self._zip_model_bundle(output_path) - print("TensorFlow model exported to", output_path) - - return v0_5.TensorflowSavedModelBundleWeightsDescr( - source=output_path, - parent="keras_hdf5", - tensorflow_version=Version(tensorflow.__version__), - ) - - # adapted from - # https://github.com/deepimagej/pydeepimagej/blob/master/pydeepimagej/yaml/create_config.py#L236 - def _convert_tf1( - self, - keras_weight_path: Path, - output_path: Path, - 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 tensorflow import ( - keras, # type: ignore - ) - - except Exception: - # if the above fails try to export with the standalone keras - import keras - - @no_type_check - def build_tf_model(): - keras_model = keras.models.load_model(keras_weight_path) - assert tensorflow is not None - builder = tensorflow.saved_model.builder.SavedModelBuilder(output_path) - signature = ( - tensorflow.saved_model.signature_def_utils.predict_signature_def( - inputs={input_name: keras_model.input}, - outputs={output_name: keras_model.output}, - ) - ) - - signature_def_map = { - tensorflow.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature - } - - builder.add_meta_graph_and_variables( - keras.backend.get_session(), - [tensorflow.saved_model.tag_constants.SERVING], - signature_def_map=signature_def_map, - ) - builder.save() - - build_tf_model() - - if zip_weights: - output_path = self._zip_model_bundle(output_path) - print("TensorFlow model exported to", output_path) - - return v0_5.TensorflowSavedModelBundleWeightsDescr( - source=output_path, - parent="keras_hdf5", - tensorflow_version=Version(tensorflow.__version__), - ) - - def _zip_model_bundle(self, model_bundle_folder: Path): - zipped_model_bundle = model_bundle_folder.with_suffix(".zip") - - with ZipFile(zipped_model_bundle, "w") as zip_obj: - for root, _, files in os.walk(model_bundle_folder): - for filename in files: - src = os.path.join(root, filename) - zip_obj.write(src, os.path.relpath(src, model_bundle_folder)) - - try: - shutil.rmtree(model_bundle_folder) - except Exception: - print("TensorFlow bundled model was not removed after compression") - - return zipped_model_bundle diff --git a/bioimageio/core/weight_converters/keras_to_tensorflow.py b/bioimageio/core/weight_converters/keras_to_tensorflow.py new file mode 100644 index 00000000..083bae5b --- /dev/null +++ b/bioimageio/core/weight_converters/keras_to_tensorflow.py @@ -0,0 +1,184 @@ +import os +import shutil +from pathlib import Path +from typing import Union, no_type_check +from zipfile import ZipFile + +import tensorflow + +from bioimageio.core.io import ensure_unzipped +from bioimageio.spec._internal.io_utils import download +from bioimageio.spec._internal.version_type import Version +from bioimageio.spec.common import ZipPath +from bioimageio.spec.model import v0_4, v0_5 + +try: + # try to build the tf model with the keras import from tensorflow + from tensorflow import keras +except Exception: + # if the above fails try to export with the standalone keras + import keras + + +def convert( + model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], *, output_path: Path +) -> v0_5.TensorflowSavedModelBundleWeightsDescr: + """ + Convert model weights from the 'keras_hdf5' format to the 'tensorflow_saved_model_bundle' format. + + This method handles the conversion of Keras HDF5 model weights into a TensorFlow SavedModel bundle, + which is the recommended format for deploying TensorFlow models. The method supports both TensorFlow 1.x + and 2.x versions, with appropriate checks to ensure compatibility. + + Adapted from: + https://github.com/deepimagej/pydeepimagej/blob/5aaf0e71f9b04df591d5ca596f0af633a7e024f5/pydeepimagej/yaml/create_config.py + + Args: + model_descr (Union[v0_4.ModelDescr, v0_5.ModelDescr]): + The bioimage.io model description containing the model's metadata and weights. + output_path (Path): + The directory where the TensorFlow SavedModel bundle will be saved. + This path must not already exist and, if necessary, will be zipped into a .zip file. + use_tracing (bool): + Placeholder argument; currently not used in this method but required to match the abstract method signature. + + Raises: + ValueError: + - If the specified `output_path` already exists. + - If the Keras HDF5 weights are missing in the model description. + RuntimeError: + If there is a mismatch between the TensorFlow version used by the model and the version installed. + NotImplementedError: + If the model has multiple inputs or outputs and TensorFlow 1.x is being used. + + Returns: + v0_5.TensorflowSavedModelBundleWeightsDescr: + A descriptor object containing information about the converted TensorFlow SavedModel bundle. + """ + tf_major_ver = int(tensorflow.__version__.split(".")[0]) + + if output_path.suffix == ".zip": + output_path = output_path.with_suffix("") + zip_weights = True + else: + zip_weights = False + + if output_path.exists(): + raise ValueError(f"The ouptut directory at {output_path} must not exist.") + + if model_descr.weights.keras_hdf5 is None: + raise ValueError("Missing Keras Hdf5 weights to convert from.") + + weight_spec = model_descr.weights.keras_hdf5 + weight_path = download(weight_spec.source).path + + if weight_spec.tensorflow_version: + model_tf_major_ver = int(weight_spec.tensorflow_version.major) + if model_tf_major_ver != tf_major_ver: + raise RuntimeError( + f"Tensorflow major versions of model {model_tf_major_ver} is not {tf_major_ver}" + ) + + if tf_major_ver == 1: + if len(model_descr.inputs) != 1 or len(model_descr.outputs) != 1: + raise NotImplementedError( + "Weight conversion for models with multiple inputs or outputs is not yet implemented." + ) + + input_name = str( + d.id + if isinstance((d := model_descr.inputs[0]), v0_5.InputTensorDescr) + else d.name + ) + output_name = str( + d.id + if isinstance((d := model_descr.outputs[0]), v0_5.OutputTensorDescr) + else d.name + ) + return _convert_tf1( + ensure_unzipped(weight_path, Path("bioimageio_unzipped_tf_weights")), + output_path, + input_name, + output_name, + zip_weights, + ) + else: + return _convert_tf2(weight_path, output_path, zip_weights) + + +def _convert_tf2( + keras_weight_path: Union[Path, ZipPath], output_path: Path, zip_weights: bool +) -> v0_5.TensorflowSavedModelBundleWeightsDescr: + model = keras.models.load_model(keras_weight_path) # type: ignore + keras.models.save_model(model, output_path) # type: ignore + + if zip_weights: + output_path = _zip_model_bundle(output_path) + print("TensorFlow model exported to", output_path) + + return v0_5.TensorflowSavedModelBundleWeightsDescr( + source=output_path, + parent="keras_hdf5", + tensorflow_version=Version(tensorflow.__version__), + ) + + +# adapted from +# https://github.com/deepimagej/pydeepimagej/blob/master/pydeepimagej/yaml/create_config.py#L236 +def _convert_tf1( + keras_weight_path: Path, + output_path: Path, + input_name: str, + output_name: str, + zip_weights: bool, +) -> v0_5.TensorflowSavedModelBundleWeightsDescr: + + @no_type_check + def build_tf_model(): + keras_model = keras.models.load_model(keras_weight_path) + assert tensorflow is not None + builder = tensorflow.saved_model.builder.SavedModelBuilder(output_path) + signature = tensorflow.saved_model.signature_def_utils.predict_signature_def( + inputs={input_name: keras_model.input}, + outputs={output_name: keras_model.output}, + ) + + signature_def_map = { + tensorflow.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature + } + + builder.add_meta_graph_and_variables( + keras.backend.get_session(), + [tensorflow.saved_model.tag_constants.SERVING], + signature_def_map=signature_def_map, + ) + builder.save() + + build_tf_model() + + if zip_weights: + output_path = _zip_model_bundle(output_path) + print("TensorFlow model exported to", output_path) + + return v0_5.TensorflowSavedModelBundleWeightsDescr( + source=output_path, + parent="keras_hdf5", + tensorflow_version=Version(tensorflow.__version__), + ) + + +def _zip_model_bundle(model_bundle_folder: Path): + zipped_model_bundle = model_bundle_folder.with_suffix(".zip") + + with ZipFile(zipped_model_bundle, "w") as zip_obj: + for root, _, files in os.walk(model_bundle_folder): + for filename in files: + src = os.path.join(root, filename) + zip_obj.write(src, os.path.relpath(src, model_bundle_folder)) + + try: + shutil.rmtree(model_bundle_folder) + except Exception: + print("TensorFlow bundled model was not removed after compression") + + return zipped_model_bundle diff --git a/bioimageio/core/weight_converters/pytorch_to_onnx.py b/bioimageio/core/weight_converters/pytorch_to_onnx.py index acb621e2..9f2b2e6f 100644 --- a/bioimageio/core/weight_converters/pytorch_to_onnx.py +++ b/bioimageio/core/weight_converters/pytorch_to_onnx.py @@ -1,28 +1,19 @@ -import abc -import os -import shutil from pathlib import Path -from typing import Any, List, Sequence, Union, cast, no_type_check -from zipfile import ZipFile +from typing import Any, List, Sequence, Union, cast import numpy as np import torch from numpy.testing import assert_array_almost_equal -from torch.jit import ScriptModule -from typing_extensions import assert_never +from bioimageio.core.backends.pytorch_backend import load_torch_model from bioimageio.core.digest_spec import get_member_id, get_test_inputs -from bioimageio.core.model_adapters._pytorch_model_adapter import PytorchModelAdapter -from bioimageio.spec._internal.io_utils import download -from bioimageio.spec._internal.version_type import Version from bioimageio.spec.model import v0_4, v0_5 -from bioimageio.spec.model.v0_5 import WeightsEntryDescrBase def convert( model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], *, - # output_path: Path, + output_path: Path, use_tracing: bool = True, test_decimal: int = 4, verbose: bool = False, @@ -81,7 +72,7 @@ def convert( ] if use_tracing: - torch.onnx.export( + _ = torch.onnx.export( model, (tuple(input_tensors) if len(input_tensors) > 1 else input_tensors[0]), str(output_path), diff --git a/bioimageio/core/weight_converters/pytorch_to_torchscript.py b/bioimageio/core/weight_converters/pytorch_to_torchscript.py new file mode 100644 index 00000000..a724e5f8 --- /dev/null +++ b/bioimageio/core/weight_converters/pytorch_to_torchscript.py @@ -0,0 +1,154 @@ +from pathlib import Path +from typing import Any, List, Sequence, Tuple, Union + +import numpy as np +import torch +from numpy.testing import assert_array_almost_equal +from torch.jit import ScriptModule +from typing_extensions import assert_never + +from bioimageio.core.backends.pytorch_backend import load_torch_model +from bioimageio.spec._internal.version_type import Version +from bioimageio.spec.model import v0_4, v0_5 + + +def convert( + model_descr: Union[v0_4.ModelDescr, v0_5.ModelDescr], + *, + output_path: Path, + use_tracing: bool = True, +) -> v0_5.TorchscriptWeightsDescr: + """ + Convert model weights from the PyTorch `state_dict` format to TorchScript. + + Args: + model_descr (Union[v0_4.ModelDescr, v0_5.ModelDescr]): + The model description object that contains the model and its weights in the PyTorch `state_dict` format. + output_path (Path): + The file path where the TorchScript model will be saved. + use_tracing (bool): + Whether to use tracing or scripting to export the TorchScript format. + - `True`: Use tracing, which is recommended for models with straightforward control flow. + - `False`: Use scripting, which is better for models with dynamic control flow (e.g., loops, conditionals). + + Raises: + ValueError: + If the provided model does not have weights in the PyTorch `state_dict` format. + + Returns: + v0_5.TorchscriptWeightsDescr: + A descriptor object that contains information about the exported TorchScript weights. + """ + state_dict_weights_descr = model_descr.weights.pytorch_state_dict + if state_dict_weights_descr is None: + raise ValueError( + "The provided model does not have weights in the pytorch state dict format" + ) + + input_data = model_descr.get_input_test_arrays() + + 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) + scripted_module: Union[ # pyright: ignore[reportUnknownVariableType] + ScriptModule, Tuple[Any, ...] + ] = ( + torch.jit.trace(model, input_data) + if use_tracing + else torch.jit.script(model) + ) + assert not isinstance(scripted_module, tuple), scripted_module + _check_predictions( + model=model, + scripted_model=scripted_module, + model_spec=model_descr, + input_data=input_data, + ) + + scripted_module.save(str(output_path)) + + return v0_5.TorchscriptWeightsDescr( + source=output_path, + pytorch_version=Version(torch.__version__), + parent="pytorch_state_dict", + ) + + +def _check_predictions( + model: Any, + scripted_model: Any, + model_spec: v0_4.ModelDescr | v0_5.ModelDescr, + input_data: Sequence[torch.Tensor], +): + def _check(input_: Sequence[torch.Tensor]) -> None: + expected_tensors = model(*input_) + if isinstance(expected_tensors, torch.Tensor): + expected_tensors = [expected_tensors] + expected_outputs: List[np.ndarray[Any, Any]] = [ + out.numpy() for out in expected_tensors + ] + + output_tensors = scripted_model(*input_) + if isinstance(output_tensors, torch.Tensor): + output_tensors = [output_tensors] + outputs: List[np.ndarray[Any, Any]] = [out.numpy() for out in output_tensors] + + try: + for exp, out in zip(expected_outputs, outputs): + assert_array_almost_equal(exp, out, decimal=4) + except AssertionError as e: + raise ValueError( + f"Results before and after weights conversion do not agree:\n {str(e)}" + ) + + _check(input_data) + + if len(model_spec.inputs) > 1: + return # FIXME: why don't we check multiple inputs? + + input_descr = model_spec.inputs[0] + if isinstance(input_descr, v0_4.InputTensorDescr): + if not isinstance(input_descr.shape, v0_4.ParameterizedInputShape): + return + min_shape = input_descr.shape.min + step = input_descr.shape.step + else: + min_shape: List[int] = [] + step: List[int] = [] + for axis in input_descr.axes: + if isinstance(axis.size, v0_5.ParameterizedSize): + min_shape.append(axis.size.min) + step.append(axis.size.step) + elif isinstance(axis.size, int): + min_shape.append(axis.size) + step.append(0) + elif axis.size is None: + raise NotImplementedError( + f"Can't verify inputs that don't specify their shape fully: {axis}" + ) + elif isinstance(axis.size, v0_5.SizeReference): + raise NotImplementedError(f"Can't handle axes like '{axis}' yet") + else: + assert_never(axis.size) + + input_tensor = input_data[0] + max_shape = input_tensor.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 step_dim == 0 + else slice(0, max_dim - step_factor * step_dim, 1) + ) + for max_dim, step_dim in zip(max_shape, step) + ) + sliced_input = input_tensor[slice_] + if any( + sliced_dim < min_dim + for sliced_dim, min_dim in zip(sliced_input.shape, min_shape) + ): + return + _check([sliced_input]) diff --git a/tests/conftest.py b/tests/conftest.py index 253ade2f..32a2b6a8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,7 +14,7 @@ import torch torch_version = tuple(map(int, torch.__version__.split(".")[:2])) - logger.warning(f"detected torch version {torch_version}.x") + logger.warning(f"detected torch version {torch.__version__}") except ImportError: torch = None torch_version = None @@ -29,7 +29,7 @@ try: import tensorflow # type: ignore - tf_major_version = int(tensorflow.__version__.split(".")[0]) # type: ignore + tf_major_version = int(tensorflow.__version__.split(".")[0]) except ImportError: tensorflow = None tf_major_version = None diff --git a/tests/test_weight_converters.py b/tests/test_weight_converters.py index 88010744..71f6ccfa 100644 --- a/tests/test_weight_converters.py +++ b/tests/test_weight_converters.py @@ -1,37 +1,36 @@ # type: ignore # TODO enable type checking +import os import zipfile from pathlib import Path import pytest -import os - -from bioimageio.spec import load_description -from bioimageio.spec.model import v0_5 - from bioimageio.core.weight_converters import ( - Pytorch2Torchscipt, Pytorch2Onnx, Tensorflow2Bundled, ) +from bioimageio.spec import load_description +from bioimageio.spec.model import v0_5 def test_torchscript_converter(any_torch_model, tmp_path): + from bioimageio.core.weight_converters.pytorch_to_torchscript import convert + bio_model = load_description(any_torch_model) out_path = tmp_path / "weights.pt" - util = Pytorch2Torchscipt() - ret_val = util.convert(bio_model, out_path) + ret_val = convert(bio_model, out_path) assert out_path.exists() assert isinstance(ret_val, v0_5.TorchscriptWeightsDescr) assert ret_val.source == out_path def test_onnx_converter(convert_to_onnx, tmp_path): + from bioimageio.core.weight_converters.pytorch_to_onnx import convert + bio_model = load_description(convert_to_onnx) out_path = tmp_path / "weights.onnx" opset_version = 15 - util = Pytorch2Onnx() - ret_val = util.convert( + ret_val = convert( model_descr=bio_model, output_path=out_path, test_decimal=3, diff --git a/tests/utils.py b/tests/utils.py index 3a8e695b..805eecfa 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,3 +1,5 @@ +"""utils to test bioimageio.core""" + import os from functools import wraps from typing import Any, Protocol, Sequence, Type From 80f9ed0286da086cc4ee10c893f01979c46ed13e Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 9 Dec 2024 14:49:15 +0100 Subject: [PATCH 28/47] update test_weight_converters.py --- tests/test_weight_converters.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/test_weight_converters.py b/tests/test_weight_converters.py index 71f6ccfa..57f40ce7 100644 --- a/tests/test_weight_converters.py +++ b/tests/test_weight_converters.py @@ -5,10 +5,7 @@ import pytest -from bioimageio.core.weight_converters import ( - Pytorch2Onnx, - Tensorflow2Bundled, -) +from bioimageio.core import test_model from bioimageio.spec import load_description from bioimageio.spec.model import v0_5 @@ -40,13 +37,17 @@ def test_onnx_converter(convert_to_onnx, tmp_path): assert isinstance(ret_val, v0_5.OnnxWeightsDescr) assert ret_val.opset_version == opset_version assert ret_val.source == out_path + bio_model.weights.onnx = ret_val + summary = test_model(bio_model, weights_format="onnx") + assert summary.status == "passed", summary.format() def test_tensorflow_converter(any_keras_model: Path, tmp_path: Path): + from bioimageio.core.weight_converters.keras_to_tensorflow import convert + model = load_description(any_keras_model) out_path = tmp_path / "weights.h5" - util = Tensorflow2Bundled() - ret_val = util.convert(model, out_path) + ret_val = convert(model, output_path=out_path) assert out_path.exists() assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr) assert ret_val.source == out_path From dad81860c1f0d5f5b4304da96b773fd7c78924d8 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 9 Dec 2024 14:49:28 +0100 Subject: [PATCH 29/47] add test_bioimageio_collection.py --- tests/test_bioimageio_collection.py | 60 +++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) create mode 100644 tests/test_bioimageio_collection.py diff --git a/tests/test_bioimageio_collection.py b/tests/test_bioimageio_collection.py new file mode 100644 index 00000000..2cf9ced0 --- /dev/null +++ b/tests/test_bioimageio_collection.py @@ -0,0 +1,60 @@ +from typing import Any, Collection, Dict, Iterable, Mapping, Tuple + +import pytest +import requests +from pydantic import HttpUrl + +from bioimageio.spec import InvalidDescr +from bioimageio.spec.common import Sha256 +from tests.utils import ParameterSet, expensive_test + +BASE_URL = "https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/" + + +def _get_latest_rdf_sources(): + entries: Any = requests.get(BASE_URL + "all_versions.json").json()["entries"] + ret: Dict[str, Tuple[HttpUrl, Sha256]] = {} + for entry in entries: + version = entry["versions"][0] + ret[f"{entry['concept']}/{version['v']}"] = ( + HttpUrl(version["source"]), # pyright: ignore[reportCallIssue] + Sha256(version["sha256"]), + ) + + return ret + + +ALL_LATEST_RDF_SOURCES: Mapping[str, Tuple[HttpUrl, Sha256]] = _get_latest_rdf_sources() + + +def yield_bioimageio_yaml_urls() -> Iterable[ParameterSet]: + for descr_url, sha in ALL_LATEST_RDF_SOURCES.values(): + key = ( + str(descr_url) + .replace(BASE_URL, "") + .replace("/files/rdf.yaml", "") + .replace("/files/bioimageio.yaml", "") + ) + yield pytest.param(descr_url, sha, key, id=key) + + +KNOWN_INVALID: Collection[str] = set() + + +@expensive_test +@pytest.mark.parametrize("descr_url,sha,key", list(yield_bioimageio_yaml_urls())) +def test_rdf( + descr_url: HttpUrl, + sha: Sha256, + key: str, +): + if key in KNOWN_INVALID: + pytest.skip("known failure") + + from bioimageio.core import load_description_and_test + + descr = load_description_and_test(descr_url, sha256=sha) + assert not isinstance(descr, InvalidDescr) + assert ( + descr.validation_summary.status == "passed" + ), descr.validation_summary.format() From e8354322f7ec6c29d89444a66128f1b35fae3dc7 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Mon, 9 Dec 2024 14:53:34 +0100 Subject: [PATCH 30/47] add onnx as dev dep --- dev/env-py38.yaml | 1 + dev/env-tf.yaml | 1 + dev/env-wo-python.yaml | 1 + dev/env.yaml | 1 + setup.py | 1 + 5 files changed, 5 insertions(+) diff --git a/dev/env-py38.yaml b/dev/env-py38.yaml index 69030cc9..23286840 100644 --- a/dev/env-py38.yaml +++ b/dev/env-py38.yaml @@ -17,6 +17,7 @@ dependencies: - loguru - matplotlib - numpy + - onnx - onnxruntime - packaging>=17.0 - pdoc diff --git a/dev/env-tf.yaml b/dev/env-tf.yaml index 799d2a59..ac443f65 100644 --- a/dev/env-tf.yaml +++ b/dev/env-tf.yaml @@ -17,6 +17,7 @@ dependencies: - loguru - matplotlib - numpy + - onnx - onnxruntime - packaging>=17.0 - pdoc diff --git a/dev/env-wo-python.yaml b/dev/env-wo-python.yaml index a0b7c978..2a77d25b 100644 --- a/dev/env-wo-python.yaml +++ b/dev/env-wo-python.yaml @@ -17,6 +17,7 @@ dependencies: - loguru - matplotlib - numpy + - onnx - onnxruntime - packaging>=17.0 - pdoc diff --git a/dev/env.yaml b/dev/env.yaml index a65158d9..7ff6abed 100644 --- a/dev/env.yaml +++ b/dev/env.yaml @@ -16,6 +16,7 @@ dependencies: - loguru - matplotlib - numpy + - onnx - onnxruntime - packaging>=17.0 - pdoc diff --git a/setup.py b/setup.py index 95414371..c1a60e40 100644 --- a/setup.py +++ b/setup.py @@ -56,6 +56,7 @@ "jupyter", "jupyter-black", "matplotlib", + "onnx", "onnxruntime", "packaging>=17.0", "pre-commit", From 459696dc50defa81966528b38c20a214a485cf60 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Tue, 10 Dec 2024 10:31:56 +0100 Subject: [PATCH 31/47] add get_pre_and_postprocessing --- bioimageio/core/proc_setup.py | 34 +++++++++++++++++++ .../core/weight_converters/pytorch_to_onnx.py | 7 +++- 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/bioimageio/core/proc_setup.py b/bioimageio/core/proc_setup.py index b9afb711..89277da5 100644 --- a/bioimageio/core/proc_setup.py +++ b/bioimageio/core/proc_setup.py @@ -1,4 +1,5 @@ from typing import ( + Callable, Iterable, List, Mapping, @@ -45,6 +46,11 @@ class PreAndPostprocessing(NamedTuple): post: List[Processing] +class _ProcessingCallables(NamedTuple): + pre: Callable[[Sample], None] + post: Callable[[Sample], None] + + class _SetupProcessing(NamedTuple): pre: List[Processing] post: List[Processing] @@ -52,6 +58,34 @@ class _SetupProcessing(NamedTuple): post_measures: Set[Measure] +class _ApplyProcs: + def __init__(self, procs: Sequence[Processing]): + super().__init__() + self._procs = procs + + def __call__(self, sample: Sample) -> None: + for op in self._procs: + op(sample) + + +def get_pre_and_postprocessing( + model: AnyModelDescr, + *, + dataset_for_initial_statistics: Iterable[Sample], + keep_updating_initial_dataset_stats: bool = False, + fixed_dataset_stats: Optional[Mapping[DatasetMeasure, MeasureValue]] = None, +) -> _ProcessingCallables: + """Creates callables to apply pre- and postprocessing in-place to a sample""" + + setup = setup_pre_and_postprocessing( + model=model, + dataset_for_initial_statistics=dataset_for_initial_statistics, + keep_updating_initial_dataset_stats=keep_updating_initial_dataset_stats, + fixed_dataset_stats=fixed_dataset_stats, + ) + return _ProcessingCallables(_ApplyProcs(setup.pre), _ApplyProcs(setup.post)) + + def setup_pre_and_postprocessing( model: AnyModelDescr, dataset_for_initial_statistics: Iterable[Sample], diff --git a/bioimageio/core/weight_converters/pytorch_to_onnx.py b/bioimageio/core/weight_converters/pytorch_to_onnx.py index 9f2b2e6f..468e56ac 100644 --- a/bioimageio/core/weight_converters/pytorch_to_onnx.py +++ b/bioimageio/core/weight_converters/pytorch_to_onnx.py @@ -7,6 +7,7 @@ from bioimageio.core.backends.pytorch_backend import load_torch_model from bioimageio.core.digest_spec import get_member_id, get_test_inputs +from bioimageio.core.proc_setup import get_pre_and_postprocessing from bioimageio.spec.model import v0_4, v0_5 @@ -61,6 +62,10 @@ def convert( input_data = [ sample.members[get_member_id(ipt)].data.data for ipt in model_descr.inputs ] + procs = get_pre_and_postprocessing( + model_descr, dataset_for_initial_statistics=[sample] + ) + procs.pre(sample) input_tensors = [torch.from_numpy(ipt) for ipt in input_data] model = load_torch_model(state_dict_weights_descr) @@ -74,7 +79,7 @@ def convert( if use_tracing: _ = torch.onnx.export( model, - (tuple(input_tensors) if len(input_tensors) > 1 else input_tensors[0]), + tuple(input_tensors), str(output_path), verbose=verbose, opset_version=opset_version, From 169cf1760447dfd0e10c79cae29cf0b4d72c391c Mon Sep 17 00:00:00 2001 From: fynnbe Date: Tue, 10 Dec 2024 10:36:30 +0100 Subject: [PATCH 32/47] use dim instead of deprecated dims arg name --- bioimageio/core/stat_calculators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bioimageio/core/stat_calculators.py b/bioimageio/core/stat_calculators.py index 41233a5b..01d553e7 100644 --- a/bioimageio/core/stat_calculators.py +++ b/bioimageio/core/stat_calculators.py @@ -137,7 +137,7 @@ def compute( else: n = int(np.prod([tensor.sizes[d] for d in self._axes])) - var = xr.dot(c, c, dims=self._axes) / n + var = xr.dot(c, c, dim=self._axes) / n assert isinstance(var, xr.DataArray) std = np.sqrt(var) assert isinstance(std, xr.DataArray) From 7d8e7fc396e899ad693ffe8a7896fde626d41bfd Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 11 Dec 2024 10:56:43 +0100 Subject: [PATCH 33/47] update tests --- tests/test_weight_converters.py | 82 ++++++++++++++++++++++++++------- 1 file changed, 66 insertions(+), 16 deletions(-) diff --git a/tests/test_weight_converters.py b/tests/test_weight_converters.py index 57f40ce7..d9e95ea8 100644 --- a/tests/test_weight_converters.py +++ b/tests/test_weight_converters.py @@ -10,25 +10,28 @@ from bioimageio.spec.model import v0_5 -def test_torchscript_converter(any_torch_model, tmp_path): +def test_pytorch_to_torchscript(any_torch_model, tmp_path): from bioimageio.core.weight_converters.pytorch_to_torchscript import convert - bio_model = load_description(any_torch_model) + model_descr = load_description(any_torch_model) out_path = tmp_path / "weights.pt" - ret_val = convert(bio_model, out_path) + ret_val = convert(model_descr, out_path) assert out_path.exists() assert isinstance(ret_val, v0_5.TorchscriptWeightsDescr) assert ret_val.source == out_path + model_descr.weights.torchscript = ret_val + summary = test_model(model_descr, weight_format="torchscript") + assert summary.status == "passed", summary.format() -def test_onnx_converter(convert_to_onnx, tmp_path): +def test_pytorch_to_onnx(convert_to_onnx, tmp_path): from bioimageio.core.weight_converters.pytorch_to_onnx import convert - bio_model = load_description(convert_to_onnx) + model_descr = load_description(convert_to_onnx) out_path = tmp_path / "weights.onnx" opset_version = 15 ret_val = convert( - model_descr=bio_model, + model_descr=model_descr, output_path=out_path, test_decimal=3, opset_version=opset_version, @@ -37,28 +40,34 @@ def test_onnx_converter(convert_to_onnx, tmp_path): assert isinstance(ret_val, v0_5.OnnxWeightsDescr) assert ret_val.opset_version == opset_version assert ret_val.source == out_path - bio_model.weights.onnx = ret_val - summary = test_model(bio_model, weights_format="onnx") + + model_descr.weights.onnx = ret_val + summary = test_model(model_descr, weight_format="onnx") assert summary.status == "passed", summary.format() -def test_tensorflow_converter(any_keras_model: Path, tmp_path: Path): +def test_keras_to_tensorflow(any_keras_model: Path, tmp_path: Path): from bioimageio.core.weight_converters.keras_to_tensorflow import convert - model = load_description(any_keras_model) - out_path = tmp_path / "weights.h5" - ret_val = convert(model, output_path=out_path) + model_descr = load_description(any_keras_model) + out_path = tmp_path / "weights" + ret_val = convert(model_descr, output_path=out_path) assert out_path.exists() assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr) assert ret_val.source == out_path + model_descr.weights.keras = ret_val + summary = test_model(model_descr, weight_format="keras_hdf5") + assert summary.status == "passed", summary.format() + @pytest.mark.skip() -def test_tensorflow_converter_zipped(any_keras_model: Path, tmp_path: Path): +def test_keras_to_tensorflow_zipped(any_keras_model: Path, tmp_path: Path): + from bioimageio.core.weight_converters.keras_to_tensorflow import convert + out_path = tmp_path / "weights.zip" - model = load_description(any_keras_model) - util = Tensorflow2Bundled() - ret_val = util.convert(model, out_path) + model_descr = load_description(any_keras_model) + ret_val = convert(model_descr, out_path) assert out_path.exists() assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr) @@ -67,3 +76,44 @@ def test_tensorflow_converter_zipped(any_keras_model: Path, tmp_path: Path): with zipfile.ZipFile(out_path, "r") as f: names = set([name for name in f.namelist()]) assert len(expected_names - names) == 0 + + model_descr.weights.keras = ret_val + summary = test_model(model_descr, weight_format="keras_hdf5") + assert summary.status == "passed", summary.format() + + +# TODO: add tensorflow_to_keras converter +# def test_tensorflow_to_keras(any_tensorflow_model: Path, tmp_path: Path): +# from bioimageio.core.weight_converters.tensorflow_to_keras import convert + +# model_descr = load_description(any_tensorflow_model) +# out_path = tmp_path / "weights.h5" +# ret_val = convert(model_descr, output_path=out_path) +# assert out_path.exists() +# assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr) +# assert ret_val.source == out_path + +# model_descr.weights.keras = ret_val +# summary = test_model(model_descr, weight_format="keras_hdf5") +# assert summary.status == "passed", summary.format() + + +# @pytest.mark.skip() +# def test_tensorflow_to_keras_zipped(any_tensorflow_model: Path, tmp_path: Path): +# from bioimageio.core.weight_converters.tensorflow_to_keras import convert + +# out_path = tmp_path / "weights.zip" +# model_descr = load_description(any_tensorflow_model) +# ret_val = convert(model_descr, out_path) + +# assert out_path.exists() +# assert isinstance(ret_val, v0_5.TensorflowSavedModelBundleWeightsDescr) + +# expected_names = {"saved_model.pb", "variables/variables.index"} +# with zipfile.ZipFile(out_path, "r") as f: +# names = set([name for name in f.namelist()]) +# assert len(expected_names - names) == 0 + +# model_descr.weights.keras = ret_val +# summary = test_model(model_descr, weight_format="keras_hdf5") +# assert summary.status == "passed", summary.format() From 8b2727e00a2a789da3c0e97676008e31263adfdd Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 11 Dec 2024 15:29:36 +0100 Subject: [PATCH 34/47] add todo --- bioimageio/core/digest_spec.py | 2 +- bioimageio/core/stat_calculators.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/bioimageio/core/digest_spec.py b/bioimageio/core/digest_spec.py index 854e6a7c..441243c6 100644 --- a/bioimageio/core/digest_spec.py +++ b/bioimageio/core/digest_spec.py @@ -339,7 +339,7 @@ def create_sample_for_model( sample_id: SampleId = None, inputs: Optional[ PerMember[Union[Tensor, xr.DataArray, NDArray[Any], Path]] - ] = None, # TODO: make non-optional + ] = None, # TODO: make non-optional # TODO: accept tuple of tensor sources **kwargs: NDArray[Any], # TODO: deprecate in favor of `inputs` ) -> Sample: """Create a sample from a single set of input(s) for a specific bioimage.io model diff --git a/bioimageio/core/stat_calculators.py b/bioimageio/core/stat_calculators.py index 01d553e7..f3aa8dcd 100644 --- a/bioimageio/core/stat_calculators.py +++ b/bioimageio/core/stat_calculators.py @@ -306,7 +306,8 @@ def _initialize(self, tensor_sizes: PerAxis[int]): out_sizes[d] = s self._dims, self._shape = zip(*out_sizes.items()) - d = int(np.prod(self._shape[1:])) # type: ignore + assert self._shape is not None + d = int(np.prod(self._shape[1:])) self._digest = [TDigest() for _ in range(d)] self._indices = product(*map(range, self._shape[1:])) From 02252acc542d8edc90ef4a08610eca14c13402ae Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 11 Dec 2024 15:30:33 +0100 Subject: [PATCH 35/47] udpate pytorch_to_onnx converter --- .../core/weight_converters/pytorch_to_onnx.py | 63 +++++++++---------- tests/test_weight_converters.py | 3 +- 2 files changed, 32 insertions(+), 34 deletions(-) diff --git a/bioimageio/core/weight_converters/pytorch_to_onnx.py b/bioimageio/core/weight_converters/pytorch_to_onnx.py index 468e56ac..71ac17a6 100644 --- a/bioimageio/core/weight_converters/pytorch_to_onnx.py +++ b/bioimageio/core/weight_converters/pytorch_to_onnx.py @@ -3,7 +3,7 @@ import numpy as np import torch -from numpy.testing import assert_array_almost_equal +from numpy.testing import assert_allclose from bioimageio.core.backends.pytorch_backend import load_torch_model from bioimageio.core.digest_spec import get_member_id, get_test_inputs @@ -16,7 +16,8 @@ def convert( *, output_path: Path, use_tracing: bool = True, - test_decimal: int = 4, + relative_tolerance: float = 1e-07, + absolute_tolerance: float = 0, verbose: bool = False, opset_version: int = 15, ) -> v0_5.OnnxWeightsDescr: @@ -31,10 +32,6 @@ def convert( The file path where the ONNX model will be saved. use_tracing (bool, optional): Whether to use tracing or scripting to export the ONNX format. Defaults to True. - test_decimal (int, optional): - The decimal precision for comparing the results between the original and converted models. - This is used in the `assert_array_almost_equal` function to check if the outputs match. - Defaults to 4. verbose (bool, optional): If True, will print out detailed information during the ONNX export process. Defaults to False. opset_version (int, optional): @@ -57,29 +54,29 @@ def convert( "The provided model does not have weights in the pytorch state dict format" ) + sample = get_test_inputs(model_descr) + procs = get_pre_and_postprocessing( + model_descr, dataset_for_initial_statistics=[sample] + ) + procs.pre(sample) + inputs_numpy = [ + sample.members[get_member_id(ipt)].data.data for ipt in model_descr.inputs + ] + inputs_torch = [torch.from_numpy(ipt) for ipt in inputs_numpy] + model = load_torch_model(state_dict_weights_descr) with torch.no_grad(): - sample = get_test_inputs(model_descr) - input_data = [ - sample.members[get_member_id(ipt)].data.data for ipt in model_descr.inputs - ] - procs = get_pre_and_postprocessing( - model_descr, dataset_for_initial_statistics=[sample] - ) - procs.pre(sample) - input_tensors = [torch.from_numpy(ipt) for ipt in input_data] - model = load_torch_model(state_dict_weights_descr) + outputs_original_torch = model(*inputs_torch) + if isinstance(outputs_original_torch, torch.Tensor): + outputs_original_torch = [outputs_original_torch] - expected_tensors = model(*input_tensors) - if isinstance(expected_tensors, torch.Tensor): - expected_tensors = [expected_tensors] - expected_outputs: List[np.ndarray[Any, Any]] = [ - out.numpy() for out in expected_tensors + outputs_original: List[np.ndarray[Any, Any]] = [ + out.numpy() for out in outputs_original_torch ] if use_tracing: _ = torch.onnx.export( model, - tuple(input_tensors), + tuple(inputs_torch), str(output_path), verbose=verbose, opset_version=opset_version, @@ -98,22 +95,24 @@ def convert( sess = rt.InferenceSession(str(output_path)) onnx_input_node_args = cast( List[Any], sess.get_inputs() - ) # fixme: remove cast, try using rt.NodeArg instead of Any - onnx_inputs = { + ) # FIXME: remove cast, try using rt.NodeArg instead of Any + inputs_onnx = { input_name.name: inp - for input_name, inp in zip(onnx_input_node_args, input_data) + for input_name, inp in zip(onnx_input_node_args, inputs_numpy) } - outputs = cast( - Sequence[np.ndarray[Any, Any]], sess.run(None, onnx_inputs) + outputs_onnx = cast( + Sequence[np.ndarray[Any, Any]], sess.run(None, inputs_onnx) ) # FIXME: remove cast try: - for exp, out in zip(expected_outputs, outputs): - assert_array_almost_equal(exp, out, decimal=test_decimal) + for out_original, out_onnx in zip(outputs_original, outputs_onnx): + assert_allclose( + out_original, out_onnx, rtol=relative_tolerance, atol=absolute_tolerance + ) except AssertionError as e: - raise ValueError( - f"Results before and after weights conversion do not agree:\n {str(e)}" - ) + raise AssertionError( + "Inference results of using original and converted weights do not match" + ) from e return v0_5.OnnxWeightsDescr( source=output_path, parent="pytorch_state_dict", opset_version=opset_version diff --git a/tests/test_weight_converters.py b/tests/test_weight_converters.py index d9e95ea8..24d2b9cb 100644 --- a/tests/test_weight_converters.py +++ b/tests/test_weight_converters.py @@ -27,13 +27,12 @@ def test_pytorch_to_torchscript(any_torch_model, tmp_path): def test_pytorch_to_onnx(convert_to_onnx, tmp_path): from bioimageio.core.weight_converters.pytorch_to_onnx import convert - model_descr = load_description(convert_to_onnx) + model_descr = load_description(convert_to_onnx, format_version="latest") out_path = tmp_path / "weights.onnx" opset_version = 15 ret_val = convert( model_descr=model_descr, output_path=out_path, - test_decimal=3, opset_version=opset_version, ) assert os.path.exists(out_path) From ab8616f2db4244af9f55e48c0434157af901cc72 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 18 Dec 2024 16:45:40 +0100 Subject: [PATCH 36/47] expose determinism to cli test command --- bioimageio/core/cli.py | 4 ++++ bioimageio/core/commands.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/bioimageio/core/cli.py b/bioimageio/core/cli.py index 1fc95310..49700b43 100644 --- a/bioimageio/core/cli.py +++ b/bioimageio/core/cli.py @@ -149,6 +149,9 @@ class TestCmd(CmdBase, WithSource): summary_path: Optional[Path] = Field(None, alias="summary-path") """Path to save validation summary as JSON file.""" + determinism: Literal["seed_only", "full"] = "seed_only" + """Modes to improve reproducibility of test outputs.""" + def run(self): sys.exit( test( @@ -158,6 +161,7 @@ def run(self): decimal=self.decimal, summary_path=self.summary_path, runtime_env=self.runtime_env, + determinism=self.determinism, ) ) diff --git a/bioimageio/core/commands.py b/bioimageio/core/commands.py index 9804a93e..92d7ddbc 100644 --- a/bioimageio/core/commands.py +++ b/bioimageio/core/commands.py @@ -30,6 +30,7 @@ def test( runtime_env: Union[ Literal["currently-active", "as-described"], Path ] = "currently-active", + determinism: Literal["seed_only", "full"] = "seed_only", ) -> int: """Test a bioimageio resource. @@ -45,6 +46,7 @@ def test( devices=[devices] if isinstance(devices, str) else devices, decimal=decimal, runtime_env=runtime_env, + determinism=determinism, ) summary.display() if summary_path is not None: From f11b42883c8eba8c993134d3fdbae58d5ba9b78b Mon Sep 17 00:00:00 2001 From: fynnbe Date: Wed, 18 Dec 2024 16:48:44 +0100 Subject: [PATCH 37/47] WIP unify model adapters --- bioimageio/core/_prediction_pipeline.py | 20 +- bioimageio/core/backends/_model_adapter.py | 129 +++++++++--- bioimageio/core/backends/onnx_backend.py | 7 +- bioimageio/core/backends/pytorch_backend.py | 81 ++++---- .../core/backends/tensorflow_backend.py | 189 +++++++++--------- .../core/backends/torchscript_backend.py | 4 - 6 files changed, 242 insertions(+), 188 deletions(-) diff --git a/bioimageio/core/_prediction_pipeline.py b/bioimageio/core/_prediction_pipeline.py index 9f5ccc3e..33fd4f33 100644 --- a/bioimageio/core/_prediction_pipeline.py +++ b/bioimageio/core/_prediction_pipeline.py @@ -121,19 +121,9 @@ def predict_sample_block( self.apply_preprocessing(sample_block) output_meta = sample_block.get_transformed_meta(self._block_transform) - output = output_meta.with_data( - { - tid: out - for tid, out in zip( - self._output_ids, - self._adapter.forward( - *(sample_block.members.get(t) for t in self._input_ids) - ), - ) - if out is not None - }, - stat=sample_block.stat, - ) + local_output = self._adapter.forward(sample_block) + + output = output_meta.with_data(local_output.members, stat=local_output.stat) if not skip_postprocessing: self.apply_postprocessing(output) @@ -157,9 +147,7 @@ def predict_sample_without_blocking( out_id: out for out_id, out in zip( self._output_ids, - self._adapter.forward( - *(sample.members.get(in_id) for in_id in self._input_ids) - ), + self._adapter.forward(sample), ) if out is not None }, diff --git a/bioimageio/core/backends/_model_adapter.py b/bioimageio/core/backends/_model_adapter.py index 99919aac..6321a8df 100644 --- a/bioimageio/core/backends/_model_adapter.py +++ b/bioimageio/core/backends/_model_adapter.py @@ -1,16 +1,36 @@ import warnings from abc import ABC, abstractmethod -from typing import List, Optional, Sequence, Tuple, Union, final +from typing import ( + Any, + List, + Literal, + Optional, + Sequence, + Tuple, + Union, + assert_never, + final, +) + +from numpy.typing import NDArray -from bioimageio.spec.model import v0_4, v0_5 +from bioimageio.core.digest_spec import get_axes_infos, get_member_ids +from bioimageio.core.sample import Sample +from bioimageio.spec.model import AnyModelDescr, v0_4, v0_5 from ..tensor import Tensor -WeightsFormat = Union[v0_4.WeightsFormat, v0_5.WeightsFormat] +SupportedWeightsFormat = Literal[ + "keras_hdf5", + "onnx", + "pytorch_state_dict", + "tensorflow_saved_model_bundle", + "torchscript", +] # Known weight formats in order of priority # First match wins -DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER: Tuple[WeightsFormat, ...] = ( +DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER: Tuple[SupportedWeightsFormat, ...] = ( "pytorch_state_dict", "tensorflow_saved_model_bundle", "torchscript", @@ -39,6 +59,22 @@ class ModelAdapter(ABC): ``` """ + def __init__(self, model_description: AnyModelDescr): + super().__init__() + self._model_descr = model_description + self._input_ids = get_member_ids(model_description.inputs) + self._output_ids = get_member_ids(model_description.outputs) + self._input_axes = [ + tuple(a.id for a in get_axes_infos(t)) for t in model_description.inputs + ] + self._output_axes = [ + tuple(a.id for a in get_axes_infos(t)) for t in model_description.outputs + ] + if isinstance(model_description, v0_4.ModelDescr): + self._input_is_optional = [False] * len(model_description.inputs) + else: + self._input_is_optional = [ipt.optional for ipt in model_description.inputs] + @final @classmethod def create( @@ -46,7 +82,7 @@ def create( model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], *, devices: Optional[Sequence[str]] = None, - weight_format_priority_order: Optional[Sequence[WeightsFormat]] = None, + weight_format_priority_order: Optional[Sequence[SupportedWeightsFormat]] = None, ): """ Creates model adapter based on the passed spec @@ -59,42 +95,44 @@ def create( ) weights = model_description.weights - errors: List[Tuple[WeightsFormat, Exception]] = [] + errors: List[Tuple[SupportedWeightsFormat, Exception]] = [] weight_format_priority_order = ( DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER if weight_format_priority_order is None else weight_format_priority_order ) # limit weight formats to the ones present - weight_format_priority_order = [ + weight_format_priority_order_present: Sequence[SupportedWeightsFormat] = [ w for w in weight_format_priority_order if getattr(weights, w) is not None ] + if not weight_format_priority_order_present: + raise ValueError( + f"None of the specified weight formats ({weight_format_priority_order}) is present ({weight_format_priority_order_present})" + ) - for wf in weight_format_priority_order: - if wf == "pytorch_state_dict" and weights.pytorch_state_dict is not None: + for wf in weight_format_priority_order_present: + if wf == "pytorch_state_dict": + assert weights.pytorch_state_dict is not None try: from .pytorch_backend import PytorchModelAdapter return PytorchModelAdapter( - outputs=model_description.outputs, - weights=weights.pytorch_state_dict, - devices=devices, + model_description=model_description, devices=devices ) except Exception as e: errors.append((wf, e)) - elif ( - wf == "tensorflow_saved_model_bundle" - and weights.tensorflow_saved_model_bundle is not None - ): + elif wf == "tensorflow_saved_model_bundle": + assert weights.tensorflow_saved_model_bundle is not None try: - from .tensorflow_backend import TensorflowModelAdapter + from .tensorflow_backend import create_tf_model_adapter - return TensorflowModelAdapter( + return create_tf_model_adapter( model_description=model_description, devices=devices ) except Exception as e: errors.append((wf, e)) - elif wf == "onnx" and weights.onnx is not None: + elif wf == "onnx": + assert weights.onnx is not None try: from .onnx_backend import ONNXModelAdapter @@ -103,7 +141,8 @@ def create( ) except Exception as e: errors.append((wf, e)) - elif wf == "torchscript" and weights.torchscript is not None: + elif wf == "torchscript": + assert weights.torchscript is not None try: from .torchscript_backend import TorchscriptModelAdapter @@ -112,7 +151,8 @@ def create( ) except Exception as e: errors.append((wf, e)) - elif wf == "keras_hdf5" and weights.keras_hdf5 is not None: + elif wf == "keras_hdf5": + assert weights.keras_hdf5 is not None # keras can either be installed as a separate package or used as part of tensorflow # we try to first import the keras model adapter using the separate package and, # if it is not available, try to load the one using tf @@ -127,6 +167,8 @@ def create( ) except Exception as e: errors.append((wf, e)) + else: + assert_never(wf) assert errors if len(weight_format_priority_order) == 1: @@ -150,12 +192,48 @@ def create( def load(self, *, devices: Optional[Sequence[str]] = None) -> None: warnings.warn("Deprecated. ModelAdapter is loaded on initialization") - @abstractmethod - def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: + def forward(self, input_sample: Sample) -> Sample: """ Run forward pass of model to get model predictions + + Note: sample id and stample stat attributes are passed through """ - # TODO: handle tensor.transpose in here and make _forward_impl the abstract impl + unexpected = [mid for mid in input_sample.members if mid not in self._input_ids] + if unexpected: + warnings.warn(f"Got unexpected input tensor IDs: {unexpected}") + + input_arrays = [ + ( + None + if (a := input_sample.members.get(in_id)) is None + else a.transpose(in_order).data.data + ) + for in_id, in_order in zip(self._input_ids, self._input_axes) + ] + output_arrays = self._forward_impl(input_arrays) + assert len(output_arrays) <= len(self._output_ids) + output_tensors = [ + None if a is None else Tensor(a, dims=d) + for a, d in zip(output_arrays, self._output_axes) + ] + return Sample( + members={ + tid: out + for tid, out in zip( + self._output_ids, + output_tensors, + ) + if out is not None + }, + stat=input_sample.stat, + id=input_sample.id, + ) + + @abstractmethod + def _forward_impl( + self, input_arrays: Sequence[Optional[NDArray[Any]]] + ) -> List[Optional[NDArray[Any]]]: + """framework specific forward implementation""" @abstractmethod def unload(self): @@ -164,5 +242,8 @@ def unload(self): The moder adapter should be considered unusable afterwards. """ + def _get_input_args_numpy(self, input_sample: Sample): + """helper to extract tensor args as transposed numpy arrays""" + create_model_adapter = ModelAdapter.create diff --git a/bioimageio/core/backends/onnx_backend.py b/bioimageio/core/backends/onnx_backend.py index 21bbcc09..8e983475 100644 --- a/bioimageio/core/backends/onnx_backend.py +++ b/bioimageio/core/backends/onnx_backend.py @@ -19,11 +19,8 @@ def __init__( model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], devices: Optional[Sequence[str]] = None, ): - super().__init__() - self._internal_output_axes = [ - tuple(a.id for a in get_axes_infos(out)) - for out in model_description.outputs - ] + super().__init__(model_description=model_description) + if model_description.weights.onnx is None: raise ValueError("No ONNX weights specified for {model_description.name}") diff --git a/bioimageio/core/backends/pytorch_backend.py b/bioimageio/core/backends/pytorch_backend.py index 74e59f30..d054ad95 100644 --- a/bioimageio/core/backends/pytorch_backend.py +++ b/bioimageio/core/backends/pytorch_backend.py @@ -3,19 +3,20 @@ from contextlib import nullcontext from io import TextIOWrapper from pathlib import Path -from typing import Any, List, Literal, Optional, Sequence, Tuple, Union +from typing import Any, List, Literal, Optional, Sequence, Union import torch from loguru import logger +from numpy.typing import NDArray from torch import nn from typing_extensions import assert_never +from bioimageio.spec._internal.type_guards import is_list, is_ndarray, is_tuple from bioimageio.spec.common import ZipPath -from bioimageio.spec.model import v0_4, v0_5 +from bioimageio.spec.model import AnyModelDescr, v0_4, v0_5 from bioimageio.spec.utils import download -from ..digest_spec import get_axes_infos, import_callable -from ..tensor import Tensor +from ..digest_spec import import_callable from ._model_adapter import ModelAdapter @@ -23,17 +24,15 @@ class PytorchModelAdapter(ModelAdapter): def __init__( self, *, - outputs: Union[ - Sequence[v0_4.OutputTensorDescr], Sequence[v0_5.OutputTensorDescr] - ], - weights: Union[ - v0_4.PytorchStateDictWeightsDescr, v0_5.PytorchStateDictWeightsDescr - ], + model_description: AnyModelDescr, devices: Optional[Sequence[Union[str, torch.device]]] = None, mode: Literal["eval", "train"] = "eval", ): - super().__init__() - self.output_dims = [tuple(a.id for a in get_axes_infos(out)) for out in outputs] + super().__init__(model_description=model_description) + weights = model_description.weights.pytorch_state_dict + if weights is None: + raise ValueError("No `pytorch_state_dict` weights found") + devices = get_devices(devices) self._model = load_torch_model(weights, load_state=True, devices=devices) if mode == "eval": @@ -46,7 +45,14 @@ def __init__( self._mode: Literal["eval", "train"] = mode self._primary_device = devices[0] - def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: + def _forward_impl( + self, input_arrays: Sequence[NDArray[Any] | None] + ) -> List[Optional[NDArray[Any]]]: + tensors = [ + None if a is None else torch.from_numpy(a).to(self._primary_device) + for a in input_arrays + ] + if self._mode == "eval": ctxt = torch.no_grad elif self._mode == "train": @@ -55,35 +61,26 @@ def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: assert_never(self._mode) with ctxt(): - tensors = [ - None if ipt is None else torch.from_numpy(ipt.data.data) - for ipt in input_tensors - ] - tensors = [ - (None if t is None else t.to(self._primary_device)) for t in tensors - ] - result: Union[Tuple[Any, ...], List[Any], Any] - result = self._model(*tensors) - if not isinstance(result, (tuple, list)): - result = [result] - - result = [ - ( - None - if r is None - else r.detach().cpu().numpy() if isinstance(r, torch.Tensor) else r - ) - for r in result # pyright: ignore[reportUnknownVariableType] - ] - if len(result) > len(self.output_dims): - raise ValueError( - f"Expected at most {len(self.output_dims)} outputs, but got {len(result)}" - ) - - return [ - None if r is None else Tensor(r, dims=out) - for r, out in zip(result, self.output_dims) - ] + model_out = self._model(*tensors) + + if is_tuple(model_out) or is_list(model_out): + model_out_seq = model_out + else: + model_out_seq = model_out = [model_out] + + result: List[Optional[NDArray[Any]]] = [] + for i, r in enumerate(model_out_seq): + if r is None: + result.append(None) + elif isinstance(r, torch.Tensor): + r_np: NDArray[Any] = r.detach().cpu().numpy() + result.append(r_np) + elif is_ndarray(r): + result.append(r) + else: + raise TypeError(f"Model output[{i}] has unexpected type {type(r)}.") + + return result def unload(self) -> None: del self._model diff --git a/bioimageio/core/backends/tensorflow_backend.py b/bioimageio/core/backends/tensorflow_backend.py index 94a8165f..37d85812 100644 --- a/bioimageio/core/backends/tensorflow_backend.py +++ b/bioimageio/core/backends/tensorflow_backend.py @@ -1,16 +1,15 @@ from pathlib import Path -from typing import List, Literal, Optional, Sequence, Union +from typing import Any, List, Literal, Optional, Sequence, Union import numpy as np import tensorflow as tf from loguru import logger +from numpy.typing import NDArray from bioimageio.core.io import ensure_unzipped from bioimageio.spec.common import FileSource -from bioimageio.spec.model import v0_4, v0_5 +from bioimageio.spec.model import AnyModelDescr, v0_4, v0_5 -from ..digest_spec import get_axes_infos -from ..tensor import Tensor from ._model_adapter import ModelAdapter @@ -29,8 +28,7 @@ def __init__( ], model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], ): - super().__init__() - self.model_description = model_description + super().__init__(model_description=model_description) tf_version = v0_5.Version(tf.__version__) model_tf_version = weights.tensorflow_version if model_tf_version is None: @@ -66,18 +64,37 @@ def __init__( weight_file = ensure_unzipped( weights.source, Path("bioimageio_unzipped_tf_weights") ) - self._network = self._get_network(weight_file) - self._internal_output_axes = [ - tuple(a.id for a in get_axes_infos(out)) - for out in model_description.outputs - ] - def _get_network( # pyright: ignore[reportUnknownParameterType] - self, weight_file: FileSource + def unload(self) -> None: + logger.warning( + "Device management is not implemented for keras yet, cannot unload model" + ) + + +class TensorflowModelAdapter(ModelAdapter): + weight_format = "tensorflow_saved_model_bundle" + + def __init__( + self, + *, + model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], + devices: Optional[Sequence[str]] = None, ): + super().__init__(model_description=model_description) + if devices is not None: + logger.warning( + f"Device management is not implemented for tensorflow yet, ignoring the devices {devices}" + ) + weight_file = ensure_unzipped( - weight_file, Path("bioimageio_unzipped_tf_weights") + weights.source, Path("bioimageio_unzipped_tf_weights") ) + self._network = str(weight_file) + + def _get_network( # pyright: ignore[reportUnknownParameterType] + self, weight_file: FileSource + ): + assert tf is not None if self.use_keras_api: try: @@ -97,22 +114,14 @@ def _get_network( # pyright: ignore[reportUnknownParameterType] 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) + return # 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( # pyright: ignore[reportUnknownParameterType] - self, *input_tensors: Optional[Tensor] + def _forward_impl( # pyright: ignore[reportUnknownParameterType] + self, input_arrays: Sequence[Optional[NDArray[Any]]] ): assert tf is not None - input_keys = [ - ipt.name if isinstance(ipt, v0_4.InputTensorDescr) else ipt.id - for ipt in self.model_description.inputs - ] - output_keys = [ - out.name if isinstance(out, v0_4.OutputTensorDescr) else out.id - for out in self.model_description.outputs - ] # TODO read from spec tag = ( # pyright: ignore[reportUnknownVariableType] tf.saved_model.tag_constants.SERVING # pyright: ignore[reportAttributeAccessIssue] @@ -136,18 +145,19 @@ def _forward_tf( # pyright: ignore[reportUnknownParameterType] # get the tensors into the graph in_names = [ # pyright: ignore[reportUnknownVariableType] - signature[signature_key].inputs[key].name for key in input_keys + signature[signature_key].inputs[key].name for key in self._input_ids ] out_names = [ # pyright: ignore[reportUnknownVariableType] - signature[signature_key].outputs[key].name for key in output_keys + signature[signature_key].outputs[key].name + for key in self._output_ids ] - in_tensors = [ + in_tf_tensors = [ graph.get_tensor_by_name( name # pyright: ignore[reportUnknownArgumentType] ) for name in in_names # pyright: ignore[reportUnknownVariableType] ] - out_tensors = [ + out_tf_tensors = [ graph.get_tensor_by_name( name # pyright: ignore[reportUnknownArgumentType] ) @@ -159,15 +169,10 @@ def _forward_tf( # pyright: ignore[reportUnknownParameterType] dict( zip( out_names, # pyright: ignore[reportUnknownArgumentType] - out_tensors, - ) - ), - dict( - zip( - in_tensors, - [None if t is None else t.data for t in input_tensors], + out_tf_tensors, ) ), + dict(zip(in_tf_tensors, input_arrays)), ) # from dict to list of tensors res = [ # pyright: ignore[reportUnknownVariableType] @@ -177,14 +182,29 @@ def _forward_tf( # pyright: ignore[reportUnknownParameterType] return res # pyright: ignore[reportUnknownVariableType] - def _forward_keras( # pyright: ignore[reportUnknownParameterType] - self, *input_tensors: Optional[Tensor] + +class KerasModelAdapter(ModelAdapter): + def __init__( + self, + *, + model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], + devices: Optional[Sequence[str]] = None, + ): + if model_description.weights.tensorflow_saved_model_bundle is None: + raise ValueError("No `tensorflow_saved_model_bundle` weights found") + + super().__init__(model_description=model_description) + if devices is not None: + logger.warning( + f"Device management is not implemented for tensorflow yet, ignoring the devices {devices}" + ) + + def _forward_impl( # pyright: ignore[reportUnknownParameterType] + self, input_arrays: Sequence[Optional[NDArray[Any]]] ): - assert self.use_keras_api - assert not isinstance(self._network, str) assert tf is not None tf_tensor = [ - None if ipt is None else tf.convert_to_tensor(ipt) for ipt in input_tensors + None if ipt is None else tf.convert_to_tensor(ipt) for ipt in input_arrays ] result = self._network(*tf_tensor) # pyright: ignore[reportUnknownVariableType] @@ -201,67 +221,42 @@ def _forward_keras( # pyright: ignore[reportUnknownParameterType] for r in result # pyright: ignore[reportUnknownVariableType] ] - def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: - if self.use_keras_api: - result = self._forward_keras( # pyright: ignore[reportUnknownVariableType] - *input_tensors - ) - else: - result = self._forward_tf( # pyright: ignore[reportUnknownVariableType] - *input_tensors - ) - return [ - ( - None - if r is None - else Tensor(r, dims=axes) # pyright: ignore[reportUnknownArgumentType] - ) - for r, axes in zip( # pyright: ignore[reportUnknownVariableType] - result, # pyright: ignore[reportUnknownArgumentType] - self._internal_output_axes, - ) - ] +def create_tf_model_adapter( + model_description: AnyModelDescr, devices: Optional[Sequence[str]] +): + tf_version = v0_5.Version(tf.__version__) + weights = model_description.weights.tensorflow_saved_model_bundle + if weights is None: + raise ValueError("No `tensorflow_saved_model_bundle` weights found") - def unload(self) -> None: + model_tf_version = weights.tensorflow_version + if model_tf_version is None: logger.warning( - "Device management is not implemented for keras yet, cannot unload model" + "The model does not specify the tensorflow version." + + f"Cannot check if it is compatible with intalled tensorflow {tf_version}." ) - - -class TensorflowModelAdapter(TensorflowModelAdapterBase): - weight_format = "tensorflow_saved_model_bundle" - - def __init__( - self, - *, - model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], - devices: Optional[Sequence[str]] = None, - ): - if model_description.weights.tensorflow_saved_model_bundle is None: - raise ValueError("missing tensorflow_saved_model_bundle weights") - - super().__init__( - devices=devices, - weights=model_description.weights.tensorflow_saved_model_bundle, - model_description=model_description, + elif model_tf_version > tf_version: + logger.warning( + f"The model specifies a newer tensorflow version than installed: {model_tf_version} > {tf_version}." ) - - -class KerasModelAdapter(TensorflowModelAdapterBase): - weight_format = "keras_hdf5" - - def __init__( - self, - *, - model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], - devices: Optional[Sequence[str]] = None, + elif (model_tf_version.major, model_tf_version.minor) != ( + tf_version.major, + tf_version.minor, ): - if model_description.weights.keras_hdf5 is None: - raise ValueError("missing keras_hdf5 weights") + logger.warning( + "The tensorflow version specified by the model does not match the installed: " + + f"{model_tf_version} != {tf_version}." + ) - super().__init__( - model_description=model_description, - devices=devices, - weights=model_description.weights.keras_hdf5, + if tf_version.major <= 1: + return TensorflowModelAdapter( + model_description=model_description, devices=devices ) + else: + return KerasModelAdapter(model_description=model_description, devices=devices) + + # TODO: check how to load tf weights without unzipping + weight_file = ensure_unzipped( + weights.source, Path("bioimageio_unzipped_tf_weights") + ) diff --git a/bioimageio/core/backends/torchscript_backend.py b/bioimageio/core/backends/torchscript_backend.py index d1882180..b0419813 100644 --- a/bioimageio/core/backends/torchscript_backend.py +++ b/bioimageio/core/backends/torchscript_backend.py @@ -40,10 +40,6 @@ def __init__( self._model = torch.jit.load(weight_path) self._model.to(self.devices[0]) self._model = self._model.eval() - self._internal_output_axes = [ - tuple(a.id for a in get_axes_infos(out)) - for out in model_description.outputs - ] def forward(self, *batch: Optional[Tensor]) -> List[Optional[Tensor]]: with torch.no_grad(): From e5bbe7a76273c595ee6a726b8595b5148e3ac2ef Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 19 Dec 2024 09:19:58 +0100 Subject: [PATCH 38/47] fix TorchscriptModelAdapter --- bioimageio/core/backends/pytorch_backend.py | 5 +-- .../core/backends/torchscript_backend.py | 34 ++++++++----------- 2 files changed, 18 insertions(+), 21 deletions(-) diff --git a/bioimageio/core/backends/pytorch_backend.py b/bioimageio/core/backends/pytorch_backend.py index d054ad95..a7fecfb7 100644 --- a/bioimageio/core/backends/pytorch_backend.py +++ b/bioimageio/core/backends/pytorch_backend.py @@ -139,11 +139,12 @@ def load_torch_state_dict( state = torch.load(f, map_location=devices[0]) incompatible = model.load_state_dict(state) - if incompatible.missing_keys: + if incompatible is not None and incompatible.missing_keys: logger.warning("Missing state dict keys: {}", incompatible.missing_keys) - if incompatible.unexpected_keys: + if incompatible is not None and incompatible.unexpected_keys: logger.warning("Unexpected state dict keys: {}", incompatible.unexpected_keys) + return model diff --git a/bioimageio/core/backends/torchscript_backend.py b/bioimageio/core/backends/torchscript_backend.py index b0419813..26924e3c 100644 --- a/bioimageio/core/backends/torchscript_backend.py +++ b/bioimageio/core/backends/torchscript_backend.py @@ -3,14 +3,13 @@ from typing import Any, List, Optional, Sequence, Union import torch +from numpy.typing import NDArray -from bioimageio.spec._internal.type_guards import is_list, is_ndarray, is_tuple +from bioimageio.spec._internal.type_guards import is_list, is_tuple from bioimageio.spec.model import v0_4, v0_5 from bioimageio.spec.utils import download -from ..digest_spec import get_axes_infos from ..model_adapters import ModelAdapter -from ..tensor import Tensor class TorchscriptModelAdapter(ModelAdapter): @@ -20,7 +19,7 @@ def __init__( model_description: Union[v0_4.ModelDescr, v0_5.ModelDescr], devices: Optional[Sequence[str]] = None, ): - super().__init__() + super().__init__(model_description=model_description) if model_description.weights.torchscript is None: raise ValueError( f"No torchscript weights found for model {model_description.name}" @@ -41,33 +40,30 @@ def __init__( self._model.to(self.devices[0]) self._model = self._model.eval() - def forward(self, *batch: Optional[Tensor]) -> List[Optional[Tensor]]: + def _forward_impl( + self, input_arrays: Sequence[Optional[NDArray[Any]]] + ) -> List[Optional[NDArray[Any]]]: + with torch.no_grad(): torch_tensor = [ - None if b is None else torch.from_numpy(b.data.data).to(self.devices[0]) - for b in batch + None if a is None else torch.from_numpy(a).to(self.devices[0]) + for a in input_arrays ] - _result: Any = self._model.forward(*torch_tensor) - if is_list(_result) or is_tuple(_result): - result: Sequence[Any] = _result + output: Any = self._model.forward(*torch_tensor) + if is_list(output) or is_tuple(output): + output_seq: Sequence[Any] = output else: - result = [_result] + output_seq = [output] - result = [ + return [ ( None if r is None else r.cpu().numpy() if isinstance(r, torch.Tensor) else r ) - for r in result + for r in output_seq ] - assert len(result) == len(self._internal_output_axes) - return [ - None if r is None else Tensor(r, dims=axes) if is_ndarray(r) else r - for r, axes in zip(result, self._internal_output_axes) - ] - def unload(self) -> None: self._devices = None del self._model From 3720e85059ff1c3ac9117f49f74d03eeec3f9958 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 19 Dec 2024 10:09:12 +0100 Subject: [PATCH 39/47] update predict_sample_without_blocking --- bioimageio/core/_prediction_pipeline.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/bioimageio/core/_prediction_pipeline.py b/bioimageio/core/_prediction_pipeline.py index 33fd4f33..5ef337a0 100644 --- a/bioimageio/core/_prediction_pipeline.py +++ b/bioimageio/core/_prediction_pipeline.py @@ -142,18 +142,7 @@ def predict_sample_without_blocking( if not skip_preprocessing: self.apply_preprocessing(sample) - output = Sample( - members={ - out_id: out - for out_id, out in zip( - self._output_ids, - self._adapter.forward(sample), - ) - if out is not None - }, - stat=sample.stat, - id=sample.id, - ) + output = self._adapter.forward(sample) if not skip_postprocessing: self.apply_postprocessing(output) From 76c27e9cb3dbbcf49376d153ecc13229a113aa9e Mon Sep 17 00:00:00 2001 From: fynnbe Date: Thu, 19 Dec 2024 10:23:41 +0100 Subject: [PATCH 40/47] ensure batch and channel axes have standardized id --- bioimageio/core/axis.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/bioimageio/core/axis.py b/bioimageio/core/axis.py index 34dfa3e1..14557b80 100644 --- a/bioimageio/core/axis.py +++ b/bioimageio/core/axis.py @@ -42,6 +42,12 @@ class Axis: id: AxisId type: Literal["batch", "channel", "index", "space", "time"] + def __post_init__(self): + if self.type == "batch": + self.id = AxisId("batch") + elif self.type == "channel": + self.id = AxisId("channel") + @classmethod def create(cls, axis: AxisLike) -> Axis: if isinstance(axis, cls): From 4cbfc5a71a9f87ebc74df213e2e245e5410adefd Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 20 Dec 2024 16:00:21 +0100 Subject: [PATCH 41/47] support validation context 'raise_errors' --- README.md | 5 +++-- bioimageio/core/_resource_tests.py | 7 +++++++ tests/test_resource_tests.py | 11 +++-------- 3 files changed, 13 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 2233ec97..9985207c 100644 --- a/README.md +++ b/README.md @@ -377,8 +377,9 @@ The model specification and its validation tools can be found at Date: Fri, 20 Dec 2024 16:25:16 +0100 Subject: [PATCH 42/47] fix ONNXModelAdapter --- bioimageio/core/backends/onnx_backend.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/bioimageio/core/backends/onnx_backend.py b/bioimageio/core/backends/onnx_backend.py index 8e983475..858b4cc1 100644 --- a/bioimageio/core/backends/onnx_backend.py +++ b/bioimageio/core/backends/onnx_backend.py @@ -2,6 +2,7 @@ from typing import Any, List, Optional, Sequence, Union import onnxruntime as rt +from numpy.typing import NDArray from bioimageio.spec._internal.type_guards import is_list, is_tuple from bioimageio.spec.model import v0_4, v0_5 @@ -35,9 +36,9 @@ def __init__( f"Device management is not implemented for onnx yet, ignoring the devices {devices}" ) - def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: - assert len(input_tensors) == len(self._input_names) - input_arrays = [None if ipt is None else ipt.data.data for ipt in input_tensors] + def _forward_impl( + self, input_arrays: Sequence[Optional[NDArray[Any]]] + ) -> List[Optional[NDArray[Any]]]: result: Any = self._session.run( None, dict(zip(self._input_names, input_arrays)) ) @@ -46,10 +47,7 @@ def forward(self, *input_tensors: Optional[Tensor]) -> List[Optional[Tensor]]: else: result_seq = [result] - return [ - None if r is None else Tensor(r, dims=axes) - for r, axes in zip(result_seq, self._internal_output_axes) - ] + return result_seq # pyright: ignore[reportReturnType] def unload(self) -> None: warnings.warn( From 3b514f881d4e49209fb6098095286e215b74e632 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 20 Dec 2024 16:34:10 +0100 Subject: [PATCH 43/47] _get_axis_type ->_guess_axis_type user provided AxisIds might be uninterpretable --- bioimageio/core/axis.py | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/bioimageio/core/axis.py b/bioimageio/core/axis.py index 14557b80..30c0d281 100644 --- a/bioimageio/core/axis.py +++ b/bioimageio/core/axis.py @@ -8,19 +8,26 @@ from bioimageio.spec.model import v0_5 -def _get_axis_type(a: Literal["b", "t", "i", "c", "x", "y", "z"]): - if a == "b": +def _guess_axis_type(a: str): + if a in ("b", "batch"): return "batch" - elif a == "t": + elif a in ("t", "time"): return "time" - elif a == "i": + elif a in ("i", "index"): return "index" - elif a == "c": + elif a in ("c", "channel"): return "channel" elif a in ("x", "y", "z"): return "space" else: - return "index" # return most unspecific axis + raise ValueError( + f"Failed to infer axis type for axis id '{a}'." + + " Consider using one of: '" + + "', '".join( + ["b", "batch", "t", "time", "i", "index", "c", "channel", "x", "y", "z"] + ) + + "'. Or creating an `Axis` object instead." + ) S = TypeVar("S", bound=str) @@ -54,10 +61,10 @@ def create(cls, axis: AxisLike) -> Axis: return axis elif isinstance(axis, Axis): return Axis(id=axis.id, type=axis.type) - elif isinstance(axis, str): - return Axis(id=AxisId(axis), type=_get_axis_type(axis)) elif isinstance(axis, v0_5.AxisBase): return Axis(id=AxisId(axis.id), type=axis.type) + elif isinstance(axis, str): + return Axis(id=AxisId(axis), type=_guess_axis_type(axis)) else: assert_never(axis) From b26433112cc171ff7277d427e3a32d07560d1547 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 20 Dec 2024 16:34:28 +0100 Subject: [PATCH 44/47] fix get_axes_infos --- bioimageio/core/digest_spec.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/bioimageio/core/digest_spec.py b/bioimageio/core/digest_spec.py index 441243c6..0d4800fb 100644 --- a/bioimageio/core/digest_spec.py +++ b/bioimageio/core/digest_spec.py @@ -35,7 +35,7 @@ ) from bioimageio.spec.utils import load_array -from .axis import AxisId, AxisInfo, AxisLike, PerAxis +from .axis import Axis, AxisId, AxisInfo, AxisLike, PerAxis from .block_meta import split_multiple_shapes_into_blocks from .common import Halo, MemberId, PerMember, SampleId, TotalNumberOfBlocks from .io import load_tensor @@ -104,14 +104,15 @@ def get_axes_infos( ], ) -> List[AxisInfo]: """get a unified, simplified axis representation from spec axes""" - return [ - ( - AxisInfo.create("i") - if isinstance(a, str) and a not in ("b", "i", "t", "c", "z", "y", "x") - else AxisInfo.create(a) - ) - for a in io_descr.axes - ] + ret: List[AxisInfo] = [] + for a in io_descr.axes: + if isinstance(a, v0_5.ANY_AXIS_TYPES): + ret.append(AxisInfo.create(Axis(id=a.id, type=a.type))) + else: + assert a in ("b", "i", "t", "c", "z", "y", "x") + ret.append(AxisInfo.create(a)) + + return ret def get_member_id( From 84f24fed86ba82ac4d3e61113ee7ae88c981a35d Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 20 Dec 2024 16:34:39 +0100 Subject: [PATCH 45/47] bump pyright version --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index c1a60e40..d740a8e3 100644 --- a/setup.py +++ b/setup.py @@ -62,7 +62,7 @@ "pre-commit", "pdoc", "psutil", # parallel pytest with 'pytest -n auto' - "pyright==1.1.390", + "pyright==1.1.391", "pytest-cov", "pytest-xdist", # parallel pytest "pytest", From 27ea9aa9aa483624ad1ebc28ec74c27e92081fa1 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Fri, 20 Dec 2024 16:34:52 +0100 Subject: [PATCH 46/47] add test cases --- tests/test_tensor.py | 29 +++++++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/tests/test_tensor.py b/tests/test_tensor.py index 33163077..e00efe04 100644 --- a/tests/test_tensor.py +++ b/tests/test_tensor.py @@ -1,3 +1,5 @@ +from typing import Sequence + import numpy as np import pytest import xarray as xr @@ -8,9 +10,19 @@ @pytest.mark.parametrize( "axes", - ["yx", "xy", "cyx", "yxc", "bczyx", "xyz", "xyzc", "bzyxc"], + [ + "yx", + "xy", + "cyx", + "yxc", + "bczyx", + "xyz", + "xyzc", + "bzyxc", + ("batch", "channel", "x", "y"), + ], ) -def test_transpose_tensor_2d(axes: str): +def test_transpose_tensor_2d(axes: Sequence[str]): tensor = Tensor.from_numpy(np.random.rand(256, 256), dims=None) transposed = tensor.transpose([AxisId(a) for a in axes]) @@ -19,9 +31,18 @@ def test_transpose_tensor_2d(axes: str): @pytest.mark.parametrize( "axes", - ["zyx", "cyzx", "yzixc", "bczyx", "xyz", "xyzc", "bzyxtc"], + [ + "zyx", + "cyzx", + "yzixc", + "bczyx", + "xyz", + "xyzc", + "bzyxtc", + ("batch", "channel", "x", "y", "z"), + ], ) -def test_transpose_tensor_3d(axes: str): +def test_transpose_tensor_3d(axes: Sequence[str]): tensor = Tensor.from_numpy(np.random.rand(64, 64, 64), dims=None) transposed = tensor.transpose([AxisId(a) for a in axes]) assert transposed.ndim == len(axes) From de759d52bd9bdfe3772f6913fdf6404ef62aae57 Mon Sep 17 00:00:00 2001 From: fynnbe Date: Tue, 7 Jan 2025 11:33:27 +0100 Subject: [PATCH 47/47] fix pip install with no-deps --- .github/workflows/build.yaml | 12 ++++++++++++ dev/env-py38.yaml | 2 +- dev/env-tf.yaml | 2 +- dev/env-wo-python.yaml | 2 +- dev/env.yaml | 2 +- 5 files changed, 16 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index 634820ad..3ed0f5df 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -42,6 +42,8 @@ jobs: create-args: >- python=${{ matrix.python-version }} post-cleanup: 'all' + env: + PIP_NO_DEPS: true - name: Install py3.8 environment if: matrix.python-version == '3.8' uses: mamba-org/setup-micromamba@v1 @@ -50,6 +52,8 @@ jobs: cache-environment: true environment-file: dev/env-py38.yaml post-cleanup: 'all' + env: + PIP_NO_DEPS: true - name: additional setup run: pip install --no-deps -e . - name: Get Date @@ -90,6 +94,8 @@ jobs: create-args: >- python=${{ matrix.python-version }} post-cleanup: 'all' + env: + PIP_NO_DEPS: true - name: Install py3.8 environment if: matrix.python-version == '3.8' uses: mamba-org/setup-micromamba@v1 @@ -98,6 +104,8 @@ jobs: cache-environment: true environment-file: dev/env-py38.yaml post-cleanup: 'all' + env: + PIP_NO_DEPS: true - name: additional setup spec run: | conda remove --yes --force bioimageio.spec || true # allow failure for cached env @@ -154,6 +162,8 @@ jobs: create-args: >- python=${{ matrix.python-version }} post-cleanup: 'all' + env: + PIP_NO_DEPS: true - name: additional setup spec run: | conda remove --yes --force bioimageio.spec || true # allow failure for cached env @@ -191,6 +201,8 @@ jobs: create-args: >- python=${{ matrix.python-version }} post-cleanup: 'all' + env: + PIP_NO_DEPS: true - name: additional setup run: pip install --no-deps -e . - name: Get Date diff --git a/dev/env-py38.yaml b/dev/env-py38.yaml index 23286840..8d7e7ecf 100644 --- a/dev/env-py38.yaml +++ b/dev/env-py38.yaml @@ -42,4 +42,4 @@ dependencies: - typing-extensions - xarray - pip: - - -e --no-deps .. + - -e .. diff --git a/dev/env-tf.yaml b/dev/env-tf.yaml index ac443f65..c28c01aa 100644 --- a/dev/env-tf.yaml +++ b/dev/env-tf.yaml @@ -42,4 +42,4 @@ dependencies: - typing-extensions - xarray - pip: - - -e --no-deps .. + - -e .. diff --git a/dev/env-wo-python.yaml b/dev/env-wo-python.yaml index 2a77d25b..08c8968e 100644 --- a/dev/env-wo-python.yaml +++ b/dev/env-wo-python.yaml @@ -42,4 +42,4 @@ dependencies: - typing-extensions - xarray - pip: - - -e --no-deps .. + - -e .. diff --git a/dev/env.yaml b/dev/env.yaml index 7ff6abed..ef715090 100644 --- a/dev/env.yaml +++ b/dev/env.yaml @@ -41,4 +41,4 @@ dependencies: - typing-extensions - xarray - pip: - - -e --no-deps .. + - -e ..