The current ONNX export test suite (tests/networks/test_convert_to_onnx.py) only covers 2 out of 40+ network architectures: UNet (2D) and SegResNet (3D). This means regressions in ONNX exportability for widely used architectures like DynUNet, UNETR, VNet, DenseNet, ResNet, etc. go completely undetected until a user hits them in production.
The convert_to_onnx() utility is a key part of MONAI's deployment story, and the gap between supported networks and tested networks is significant.
Describing the solution
Add parameterized ONNX export tests for additional architectures that are known to export cleanly. Starting with these 10 LOW-risk networks that have no ONNX-incompatible patterns (no .item(), no data-dependent control flow, single-tensor output):
- DynUNet — popular for nnU-Net style pipelines
- AttentionUnet — standard attention-gated U-Net
- BasicUNet — simple baseline U-Net
- BasicUNetPlusPlus — UNet++ (with
deep_supervision=False)
- VNet — classic 3D segmentation
- HighResNet — standard 3D segmentation
- DenseNet — classification backbone
- ResNet — classification backbone
- SENet — squeeze-excitation network
- UNETR — transformer-based segmentation
All tests will use small model configurations and tiny input tensors to keep CI runtime low, following the existing parameterized pattern in the test file.
Describe alternatives
- Testing all 40+ networks at once — impractical due to ONNX-incompatible patterns in some architectures (e.g., SwinUNETR uses
.item() during init, AHNet has data-dependent control flow). Better to start with the safe subset and expand incrementally.
Additional context
The existing FIXME at line 29 of test_convert_to_onnx.py notes that CUDA produces different outputs vs ONNX — all tests run CPU-only, which is fine for verifying exportability.
The current ONNX export test suite (
tests/networks/test_convert_to_onnx.py) only covers 2 out of 40+ network architectures:UNet(2D) andSegResNet(3D). This means regressions in ONNX exportability for widely used architectures like DynUNet, UNETR, VNet, DenseNet, ResNet, etc. go completely undetected until a user hits them in production.The
convert_to_onnx()utility is a key part of MONAI's deployment story, and the gap between supported networks and tested networks is significant.Describing the solution
Add parameterized ONNX export tests for additional architectures that are known to export cleanly. Starting with these 10 LOW-risk networks that have no ONNX-incompatible patterns (no
.item(), no data-dependent control flow, single-tensor output):deep_supervision=False)All tests will use small model configurations and tiny input tensors to keep CI runtime low, following the existing parameterized pattern in the test file.
Describe alternatives
.item()during init, AHNet has data-dependent control flow). Better to start with the safe subset and expand incrementally.Additional context
The existing FIXME at line 29 of
test_convert_to_onnx.pynotes that CUDA produces different outputs vs ONNX — all tests run CPU-only, which is fine for verifying exportability.