Skip to content

Expand ONNX export test coverage beyond UNet and SegResNet #9072

Description

@chhayankjain

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions