update torch export

Author: FynnBe

src/bioimageio/core/weight_converters/_utils_onnx.py

@@ -1,15 +1,49 @@
 from collections import defaultdict
-from itertools import chain
-from typing import DefaultDict, Dict
+from typing import DefaultDict, Dict, Optional, Tuple
 
-from bioimageio.spec.model.v0_5 import ModelDescr
+from bioimageio.spec.model.v0_5 import (
+    InputAxis,
+    ModelDescr,
+    ParameterizedSize,
+    SizeReference,
+)
+from torch.export import Dim
+from typing_extensions import assert_never
 
 
-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):
+def get_dynamic_shapes(model_descr: ModelDescr):
+    dynamic_shapes: DefaultDict[str, Dict[int, Optional[Dim]]] = defaultdict(dict)
+    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:
         for i, ax in enumerate(d.axes):
-            if not isinstance(ax.size, int):
-                dynamic_axes[str(d.id)][i] = str(ax.id)
+            dim_name = f"{d.id}_{ax.id}"
+            if isinstance(ax.size, int):
+                dim = None  # fixed size (could also be left out)
+            elif ax.size is None:
+                dim = Dim(dim_name, min=1)
+            elif isinstance(ax.size, ParameterizedSize):
+                dim = Dim(dim_name, min=ax.size.min)
+            elif isinstance(ax.size, SizeReference):
+                continue  # handled below
+            else:
+                assert_never(ax.size)
 
-    return dynamic_axes
+            dynamic_shapes[str(d.id)][i] = dim
+            potential_ref_axes[dim_name] = (ax, i)
+
+    # add dynamic dims from size references
+    for d in model_descr.inputs:
+        for i, ax in enumerate(d.axes):
+            if not isinstance(ax.size, SizeReference):
+                continue  # handled above
+
+            dim_name_ref = f"{ax.size.tensor_id}_{ax.size.axis_id}"
+            ax_ref, i_ref = potential_ref_axes[dim_name_ref]
+            a = ax_ref.scale / ax.scale
+            b = ax.size.offset
+            dim_ref = dynamic_shapes[str(ax.size.tensor_id)][i_ref]
+            dim = a * dim_ref + b if dim_ref is not None else None
+            dynamic_shapes[str(d.id)][i] = dim
+
+    return dict(dynamic_shapes)

src/bioimageio/core/weight_converters/pytorch_to_onnx.py

@@ -1,14 +1,13 @@
 from pathlib import Path
 
 import torch
-
 from bioimageio.spec.model.v0_5 import ModelDescr, OnnxWeightsDescr
 
 from .. import __version__
 from ..backends.pytorch_backend import load_torch_model
 from ..digest_spec import get_member_id, get_test_input_sample
 from ..proc_setup import get_pre_and_postprocessing
-from ._utils_onnx import get_dynamic_axes
+from ._utils_onnx import get_dynamic_shapes
 
 
 def convert(
@@ -64,9 +63,10 @@ def convert(
             model,
             tuple(inputs_torch),
             str(output_path),
+            dynamo=True,
             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),
+            dynamic_shapes=get_dynamic_shapes(model_descr),
             verbose=verbose,
             opset_version=opset_version,
         )

src/bioimageio/core/weight_converters/torchscript_to_onnx.py

@@ -1,13 +1,12 @@
 from pathlib import Path
 
 import torch.jit
-
 from bioimageio.spec.model.v0_5 import ModelDescr, OnnxWeightsDescr
 
 from .. import __version__
 from ..digest_spec import get_member_id, get_test_input_sample
 from ..proc_setup import get_pre_and_postprocessing
-from ._utils_onnx import get_dynamic_axes
+from ._utils_onnx import get_dynamic_shapes
 
 
 def convert(
@@ -68,9 +67,10 @@ def convert(
             model,  # type: ignore
             tuple(inputs_torch),
             str(output_path),
+            dynamo=True,
             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),
+            dynamic_shapes=get_dynamic_shapes(model_descr),
             verbose=verbose,
             opset_version=opset_version,
         )