Project-MONAI / Project-MONAI/MONAI
Expand ONNX export test coverage beyond UNet and SegResNet
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 8.7k
- Forks
- 1.6k
- Avg merge
- 5d 1h
- Merged PRs (30d)
- 20
Description
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.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with tests/networks/test_convert_to_onnx.py and follow its existing parameterized ONNX export pattern. Add CPU-only coverage for the ten named network architectures using small configurations and tiny inputs, with BasicUNetPlusPlus configured with deep_supervision=False. Done means each listed model exports successfully without substantially increasing CI runtime.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, testing-qa
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 75/100