Support for ONNX subgraphs in Burn for conditional inference computation

[antimora]

Description

Introduce the capability to support ONNX subgraphs within the Burn framework, particularly for handling conditional inference computations. This feature aims to enable Burn to fully support modern complex ONNX models that rely on subgraphs for conditional computations.

Motivation

1. Complete ONNX Compatibility: Many advanced models like Version 4 of the Silero VAD make use of subgraphs for various tasks. Lack of support for ONNX subgraphs would mean incomplete compatibility between Burn and ONNX.

2. Conditional Inference: Subgraphs are particularly useful for models that require conditional computation during inference, allowing for more dynamic and efficient operations.

3. Model Integrity and Flexibility: Providing support for ONNX subgraphs ensures that Burn can accommodate more complex model architectures without requiring modifications to the original ONNX model.

4. Elevate Burn's Capabilities: Incorporating subgraph support would make Burn a more versatile and robust framework, increasing its appeal to a broader audience.

Proposed Implementation Steps

Backend Steps

1. Subgraph Detection: Parse the ONNX file to identify and isolate subgraphs.

2. Subgraph-to-Module Conversion: Convert these subgraphs into Burn-compatible modules, maintaining their conditional inference capabilities.

3. Integration and Testing: Integrate the converted modules into Burn and conduct thorough tests to ensure they work as expected.

Code Changes

• Some refactoring of the existing implementation will be required.

• However, this feature should be achievable without major architectural changes.

Additional Features

• Documentation: Offer detailed documentation and examples that show how to work with ONNX subgraphs within Burn.

• Error Handling and Warnings: Implement robust error handling to notify users of any potential issues during the import and conversion process.

• Logging: Provide logging mechanisms for debugging and monitoring the performance of subgraph handling.

[hkrutzer]

Thanks for creating and open-sourcing Burn! I think there is a lot of demand for running Silero VAD inside Rust, there's several projects using Whisper that want to add a VAD. Looking forward to this feature!