Expand ONNX export test coverage to 12 network architectures

The ONNX export tests only covered UNet (2D) and SegResNet (3D), leaving
regressions in exportability for widely used architectures undetected.

Add parameterized ONNX export tests for 10 additional networks:
DynUNet, AttentionUnet, BasicUNet, BasicUNetPlusPlus, VNet, HighResNet,
DenseNet, ResNet (resnet10), SEResNet50, UNETR, and FullyConnectedNet.

All tests use small model configurations and tiny inputs to keep CI fast.

Fixes #9072

Signed-off-by: chhayankjain <chhayank44@gmail.com>

Author: chhayankjain

tests/networks/test_convert_to_onnx.py

@@ -19,7 +19,21 @@
 from parameterized import parameterized
 
 from monai.networks import convert_to_onnx
-from monai.networks.nets import SegResNet, UNet
+from monai.networks.nets import (
+    UNETR,
+    AttentionUnet,
+    BasicUNet,
+    BasicUNetPlusPlus,
+    DenseNet,
+    DynUNet,
+    FullyConnectedNet,
+    HighResNet,
+    SegResNet,
+    SEResNet50,
+    UNet,
+    VNet,
+    resnet10,
+)
 from tests.test_utils import SkipIfNoModule, optional_import, skip_if_quick
 
 onnx, _ = optional_import("onnx")
@@ -32,6 +46,7 @@
 
 TESTS = list(itertools.product(TORCH_DEVICE_OPTIONS, [True, False], [True, False]))
 TESTS_ORT = list(itertools.product(TORCH_DEVICE_OPTIONS, [True]))
+TESTS_TRACE = list(itertools.product(TORCH_DEVICE_OPTIONS, [True, False]))
 
 ON_AARCH64 = platform.machine() == "aarch64"
 if ON_AARCH64:
@@ -40,6 +55,21 @@
     rtol, atol = 1e-2, 1e-2
 
 
+def _check_ort_available(test_case):
+    """Skip the test if onnxruntime is not installed.
+
+    Args:
+        test_case: the ``unittest.TestCase`` instance to call ``skipTest`` on
+            when onnxruntime is unavailable.
+
+    Raises:
+        unittest.SkipTest: when onnxruntime cannot be imported.
+    """
+    _, has_onnxruntime = optional_import("onnxruntime")
+    if not has_onnxruntime:
+        test_case.skipTest("onnxruntime is not installed probably due to python version >= 3.11.")
+
+
 @SkipIfNoModule("onnx")
 @skip_if_quick
 class TestConvertToOnnx(unittest.TestCase):
@@ -103,6 +133,323 @@ def test_seg_res_net(self, device, use_ort):
         )
         self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
 
+    @parameterized.expand(TESTS_TRACE)
+    def test_dynunet(self, device, use_ort):
+        """Test converting DynUNet to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = DynUNet(
+            spatial_dims=3,
+            in_channels=1,
+            out_channels=2,
+            kernel_size=[3, 3, 3],
+            strides=[1, 2, 2],
+            upsample_kernel_size=[2, 2],
+        )
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 32, 32, 32), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_TRACE)
+    def test_attention_unet(self, device, use_ort):
+        """Test converting AttentionUnet to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = AttentionUnet(spatial_dims=3, in_channels=1, out_channels=2, channels=(16, 32, 64), strides=(2, 2))
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 32, 32, 32), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_TRACE)
+    def test_basic_unet(self, device, use_ort):
+        """Test converting BasicUNet to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = BasicUNet(spatial_dims=3, in_channels=1, out_channels=2, features=(8, 8, 16, 32, 64, 8))
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 32, 32, 32), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_TRACE)
+    def test_basic_unet_plus_plus(self, device, use_ort):
+        """Test converting BasicUNetPlusPlus to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = BasicUNetPlusPlus(
+            spatial_dims=3, in_channels=1, out_channels=2, features=(8, 8, 16, 32, 64, 8), deep_supervision=False
+        )
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 32, 32, 32), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_TRACE)
+    def test_vnet(self, device, use_ort):
+        """Test converting VNet to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = VNet(spatial_dims=3, in_channels=1, out_channels=1)
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 32, 32, 32), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_TRACE)
+    def test_highresnet(self, device, use_ort):
+        """Test converting HighResNet to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = HighResNet(spatial_dims=3, in_channels=1, out_channels=2)
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 16, 16, 16), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_TRACE)
+    def test_densenet(self, device, use_ort):
+        """Test converting DenseNet to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = DenseNet(
+            spatial_dims=3, in_channels=1, out_channels=2, init_features=16, growth_rate=8, block_config=(2, 2, 2, 2)
+        )
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 32, 32, 32), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_TRACE)
+    def test_resnet(self, device, use_ort):
+        """Test converting ResNet to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = resnet10(pretrained=False, spatial_dims=3, n_input_channels=1, num_classes=2)
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 32, 32, 32), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_TRACE)
+    def test_seresnet(self, device, use_ort):
+        """Test converting SEResNet50 to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = SEResNet50(layers=(1, 1, 1, 1), spatial_dims=3, in_channels=1, num_classes=2)
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 32, 32, 32), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_TRACE)
+    def test_unetr(self, device, use_ort):
+        """Test converting UNETR to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = UNETR(
+            in_channels=1,
+            out_channels=2,
+            img_size=(32, 32, 32),
+            feature_size=8,
+            hidden_size=128,
+            mlp_dim=256,
+            num_heads=8,
+            spatial_dims=3,
+        )
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 32, 32, 32), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_TRACE)
+    def test_fully_connected_net(self, device, use_ort):
+        """Test converting FullyConnectedNet to ONNX with and without ORT verification.
+
+        Args:
+            device: torch device string (e.g. ``"cpu"``).
+            use_ort: if ``True``, verify via onnxruntime; if ``False``, verify
+                via ``onnx.reference.ReferenceEvaluator``. Skipped when
+                onnxruntime is unavailable.
+        """
+        if use_ort:
+            _check_ort_available(self)
+        model = FullyConnectedNet(in_channels=10, out_channels=2, hidden_channels=[20, 10])
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((4, 10), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
 
 if __name__ == "__main__":
     unittest.main()