Skip to content

Expand ONNX export test coverage to 12 network architectures #9073

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chhayankjain wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

Expand ONNX export test coverage to 12 network architectures #9073

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
349 changes: 348 additions & 1 deletion tests/networks/test_convert_to_onnx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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:
Expand All @@ -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.")
Comment thread
coderabbitai[bot] marked this conversation as resolved.


@SkipIfNoModule("onnx")
@skip_if_quick
class TestConvertToOnnx(unittest.TestCase):
Expand Down Expand Up @@ -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()
Loading