Project-MONAI / Project-MONAI/MONAI

Expand ONNX export test coverage beyond UNet and SegResNet

Open
#9,072 0 comments 0 reactions 0 assignees View on GitHub

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):

  1. DynUNet — popular for nnU-Net style pipelines
  2. AttentionUnet — standard attention-gated U-Net
  3. BasicUNet — simple baseline U-Net
  4. BasicUNetPlusPlus — UNet++ (with deep_supervision=False)
  5. VNet — classic 3D segmentation
  6. HighResNet — standard 3D segmentation
  7. DenseNet — classification backbone
  8. ResNet — classification backbone
  9. SENet — squeeze-excitation network
  10. 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.