improve onnx converters

Author: FynnBe

bioimageio/core/weight_converters/_utils_onnx.py

@@ -0,0 +1,15 @@
+from collections import defaultdict
+from itertools import chain
+from typing import DefaultDict, Dict
+
+from bioimageio.spec.model.v0_5 import ModelDescr
+
+
+def get_dynamic_axes(model_descr: ModelDescr):
+    dynamic_axes: DefaultDict[str, Dict[int, str]] = defaultdict(dict)
+    for d in chain(model_descr.inputs, model_descr.outputs):
+        for i, ax in enumerate(d.axes):
+            if not isinstance(ax.size, int):
+                dynamic_axes[str(d.id)][i] = str(ax.id)
+
+    return dynamic_axes

bioimageio/core/weight_converters/pytorch_to_onnx.py

@@ -8,14 +8,15 @@
 from ..backends.pytorch_backend import load_torch_model
 from ..digest_spec import get_member_id, get_test_inputs
 from ..proc_setup import get_pre_and_postprocessing
+from ._utils_onnx import get_dynamic_axes
 
 
 def convert(
     model_descr: ModelDescr,
     output_path: Path,
     *,
     verbose: bool = False,
-    opset_version: int = 20,
+    opset_version: int = 15,
 ) -> OnnxWeightsDescr:
     """
     Convert model weights from the Torchscript state_dict format to the ONNX format.
@@ -63,6 +64,9 @@ def convert(
             model,
             tuple(inputs_torch),
             str(output_path),
+            input_names=[str(d.id) for d in model_descr.inputs],
+            output_names=[str(d.id) for d in model_descr.outputs],
+            dynamic_axes=get_dynamic_axes(model_descr),
             verbose=verbose,
             opset_version=opset_version,
         )

bioimageio/core/weight_converters/torchscript_to_onnx.py

@@ -5,8 +5,10 @@
 from bioimageio.spec.model.v0_5 import ModelDescr, OnnxWeightsDescr
 from bioimageio.spec.utils import download
 
+from .. import __version__
 from ..digest_spec import get_member_id, get_test_inputs
 from ..proc_setup import get_pre_and_postprocessing
+from ._utils_onnx import get_dynamic_axes
 
 
 def convert(
@@ -67,10 +69,16 @@ def convert(
             model,  # type: ignore
             tuple(inputs_torch),
             str(output_path),
+            input_names=[str(d.id) for d in model_descr.inputs],
+            output_names=[str(d.id) for d in model_descr.outputs],
+            dynamic_axes=get_dynamic_axes(model_descr),
             verbose=verbose,
             opset_version=opset_version,
         )
 
     return OnnxWeightsDescr(
-        source=output_path, parent="pytorch_state_dict", opset_version=opset_version
+        source=output_path,
+        parent="torchscript",
+        opset_version=opset_version,
+        comment=f"Converted with bioimageio.core {__version__}.",
     )