diff --git a/tests/networks/test_convert_to_onnx.py b/tests/networks/test_convert_to_onnx.py index 1d4e4ea385..863d52be47 100644 --- a/tests/networks/test_convert_to_onnx.py +++ b/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()