improve dynamo export and update onnx adapter

(for external weight data files)

Author: FynnBe

src/bioimageio/core/backends/onnx_backend.py

@@ -1,11 +1,14 @@
 # pyright: reportUnknownVariableType=false
+import shutil
+import tempfile
 import warnings
+from pathlib import Path
 from typing import Any, List, Optional, Sequence, Union
 
 import onnxruntime as rt  # pyright: ignore[reportMissingTypeStubs]
-from numpy.typing import NDArray
-
 from bioimageio.spec.model import v0_4, v0_5
+from loguru import logger
+from numpy.typing import NDArray
 
 from ..model_adapters import ModelAdapter
 from ..utils._type_guards import is_list, is_tuple
@@ -20,11 +23,63 @@ def __init__(
     ):
         super().__init__(model_description=model_description)
 
-        if model_description.weights.onnx is None:
+        onnx_descr = model_description.weights.onnx
+        if onnx_descr is None:
             raise ValueError("No ONNX weights specified for {model_description.name}")
 
-        reader = model_description.weights.onnx.get_reader()
-        self._session = rt.InferenceSession(reader.read())
+        providers = None
+        if hasattr(rt, "get_available_providers"):
+            providers = rt.get_available_providers()
+
+        if (
+            isinstance(onnx_descr, v0_5.OnnxWeightsDescr)
+            and onnx_descr.external_data is not None
+        ):
+            src = onnx_descr.source.absolute()
+            src_data = onnx_descr.external_data.source.absolute()
+            if (
+                isinstance(src, Path)
+                and isinstance(src_data, Path)
+                and src.parent == src_data.parent
+            ):
+                logger.debug(
+                    "Loading ONNX model with external data from {}",
+                    src.parent,
+                )
+                self._session = rt.InferenceSession(
+                    src,
+                    providers=providers,  # pyright: ignore[reportUnknownArgumentType]
+                )
+            else:
+                src_reader = onnx_descr.get_reader()
+                src_data_reader = onnx_descr.external_data.get_reader()
+                with tempfile.TemporaryDirectory() as tmpdir:
+                    logger.debug(
+                        "Loading ONNX model with external data from {}",
+                        tmpdir,
+                    )
+                    src = Path(tmpdir) / src_reader.original_file_name
+                    src_data = Path(tmpdir) / src_data_reader.original_file_name
+                    with src.open("wb") as f:
+                        shutil.copyfileobj(src_reader, f)
+                    with src_data.open("wb") as f:
+                        shutil.copyfileobj(src_data_reader, f)
+
+                    self._session = rt.InferenceSession(
+                        src,
+                        providers=providers,  # pyright: ignore[reportUnknownArgumentType]
+                    )
+        else:
+            # load single source file from bytes (without external data, so probably <2GB)
+            logger.debug(
+                "Loading ONNX model from bytes (read from {})", onnx_descr.source
+            )
+            reader = onnx_descr.get_reader()
+            self._session = rt.InferenceSession(
+                reader.read(),
+                providers=providers,  # pyright: ignore[reportUnknownArgumentType]
+            )
+
         onnx_inputs = self._session.get_inputs()
         self._input_names: List[str] = [ipt.name for ipt in onnx_inputs]
 

src/bioimageio/core/weight_converters/_utils_torch_onnx.py

@@ -7,6 +7,7 @@
 
 import torch
 from bioimageio.spec.model.v0_5 import (
+    BatchAxis,
     FileDescr,
     InputAxis,
     ModelDescr,
@@ -71,17 +72,17 @@ def _get_dynamic_shapes_impl(model_descr: ModelDescr):
         # dynamic shapes as list to match the source code which may have
         # different arg names than the tensor ids in the model description
 
-        dynamic_shapes: List[Dict[int, TensorDim]] = []
+        dynamic_shapes: List[Dict[int, Union[int, TensorDim]]] = []
         potential_ref_axes: Dict[str, Tuple[InputAxis, int]] = {}
         # add dynamic dims from parameterized input sizes (and fixed sizes as None)
         for d in model_descr.inputs:
-            dynamic_tensor_dims: Dict[int, TensorDim] = {}
+            dynamic_tensor_dims: Dict[int, Union[int, TensorDim]] = {}
             for i, ax in enumerate(d.axes):
                 dim_name = f"{d.id}_{ax.id}"
                 if isinstance(ax.size, int):
-                    dim = STATIC_DIM  # fixed size
-                elif ax.size is None:
-                    dim = Dim(dim_name, min=1)
+                    dim = ax.size
+                elif isinstance(ax, BatchAxis):
+                    dim = Dim("batch", min=1)
                 elif isinstance(ax.size, ParameterizedSize):
                     dim = Dim(dim_name, min=ax.size.min)
                 elif isinstance(ax.size, SizeReference):

tests/test_weight_converters.py

@@ -4,16 +4,15 @@
 from pathlib import Path
 
 import pytest
-
-from bioimageio.spec import load_description
+from bioimageio.spec import load_model_description
 from bioimageio.spec.model import v0_5
 
 
 def test_pytorch_to_torchscript(any_torch_model, tmp_path):
     from bioimageio.core import test_model
     from bioimageio.core.weight_converters.pytorch_to_torchscript import convert
 
-    model_descr = load_description(any_torch_model, perform_io_checks=False)
+    model_descr = load_model_description(any_torch_model, perform_io_checks=False)
     if model_descr.implemented_format_version_tuple[:2] == (0, 4):
         pytest.skip("cannot convert to old 0.4 format")
 
@@ -31,9 +30,9 @@ def test_pytorch_to_onnx(convert_to_onnx, tmp_path):
     from bioimageio.core import test_model
     from bioimageio.core.weight_converters.pytorch_to_onnx import convert
 
-    model_descr = load_description(convert_to_onnx, format_version="latest")
+    model_descr = load_model_description(convert_to_onnx, format_version="latest")
     out_path = tmp_path / "weights.onnx"
-    opset_version = 15
+    opset_version = 18
     ret_val = convert(
         model_descr=model_descr,
         output_path=out_path,
@@ -55,7 +54,7 @@ def test_keras_to_tensorflow(any_keras_model: Path, tmp_path: Path):
     from bioimageio.core.weight_converters.keras_to_tensorflow import convert
 
     out_path = tmp_path / "weights.zip"
-    model_descr = load_description(any_keras_model)
+    model_descr = load_model_description(any_keras_model)
     ret_val = convert(model_descr, out_path)
 
     assert out_path.exists()
@@ -75,7 +74,7 @@ def test_keras_to_tensorflow(any_keras_model: Path, tmp_path: Path):
 # 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)
+#     model_descr = load_model_description(any_tensorflow_model)
 #     out_path = tmp_path / "weights.h5"
 #     ret_val = convert(model_descr, output_path=out_path)
 #     assert out_path.exists()
@@ -92,7 +91,7 @@ def test_keras_to_tensorflow(any_keras_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)
+#     model_descr = load_model_description(any_tensorflow_model)
 #     ret_val = convert(model_descr, out_path)
 
 #     assert out_path.exists()