ONNX Manager Implementation Summary

See Also:

Overview

This implementation adds support for ONNX (Open Neural Network Exchange) model evaluation to RDFAnalyzerCore, similar to the existing BDTManager functionality. ONNX is an open standard for machine learning models that enables interoperability between different ML frameworks.

What Was Implemented

1. ONNX Runtime Integration (CMake)

2. OnnxManager Plugin

Features:

Key Methods:

3. Comprehensive Unit Tests

Test Coverage:

4. Documentation

Additions:

Configuration Format

Create a configuration file (e.g., cfg/onnx_models.txt):

file=path/to/model.onnx name=model_output inputVariables=var1,var2,var3 runVar=should_run

Add to main config:

onnxConfig=cfg/onnx_models.txt

Input Padding for Fixed-Size Models

ONNX models with fixed-size inputs (e.g., transformer attention mechanisms) now support automatic zero-padding. This is particularly useful for models that require a specific input size but process variable-length sequences.

Configuration:

Add paddingSize parameter to your model configuration:

# onnx_models.txt
file=transformer.onnx name=ParT inputVariables=pt,eta,phi paddingSize=128

Behavior:

Example:

// Configuration: paddingSize=128
// Input: RVec<float>{1.0, 2.0, 3.0}  (size 3)
// After padding: RVec<float>{1.0, 2.0, 3.0, 0.0, 0.0, ..., 0.0}  (size 128)

// Input: RVec<float>{...300 elements...}  (size 300)
// After truncation: RVec<float>{...first 128 elements...}  (size 128)

This feature enables the use of transformer-based models (ParticleNet, ParT, etc.) that require fixed-size attention masks and positional encodings.

Usage Example

Single Output Model

// Create and configure OnnxManager
auto onnxManager = std::make_unique<OnnxManager>(*configProvider);
ManagerContext ctx{*configProvider, *dataManager, *systematicManager, *logger, *skimSink, *metaSink};
onnxManager->setContext(ctx);

// Define input features first
dataManager->Define("pt", ...);
dataManager->Define("eta", ...);

// Then apply models
onnxManager->applyAllModels();

// Access results in dataframe
auto df = dataManager->getDataFrame();
auto predictions = df.Take<float>("model_output");

Multi-Output Model (ParticleTransformer Style)

// Define inputs
dataManager->Define("jet_pt", ...);
dataManager->Define("jet_eta", ...);
dataManager->Define("jet_phi", ...);

// Apply multi-output model
onnxManager->applyModel("particle_transformer");

// Access individual outputs
auto df = dataManager->getDataFrame();
auto output0 = df.Take<float>("particle_transformer_output0");  // First output
auto output1 = df.Take<float>("particle_transformer_output1");  // Second output
auto output2 = df.Take<float>("particle_transformer_output2");  // Third output

Testing Notes

Build Requirements

Running Tests

The implementation includes comprehensive unit tests that can be run with:

source env.sh
source build.sh
cd build
ctest -R OnnxManagerTest

Note: Tests require a ROOT environment with CVMFS or local ROOT installation.

Design Decisions

  1. Binary Distribution: ONNX Runtime is downloaded as pre-built binaries rather than built from source to:
    • Avoid long build times
    • Reduce build complexity
    • Ensure consistent behavior across platforms
    • Follow the problem statement requirement
  2. Conditional Execution: Like BDTManager, models return -1.0 when runVar is false to:
    • Skip expensive inference when not needed
    • Maintain compatibility with existing analysis patterns
    • Enable clear identification of skipped events
  3. Code Structure: Follows BDTManager patterns to:
    • Maintain consistency across the codebase
    • Leverage existing NamedObjectManager infrastructure
    • Make the API familiar to existing users
  4. Thread Safety: Configured with single-threaded ONNX inference per event to:
    • Allow ROOT’s ImplicitMT to handle parallelism at the event level
    • Avoid thread contention within ONNX Runtime
    • Maintain consistent behavior with other managers
  5. Deferred Execution: Models are loaded during construction but NOT applied automatically to:
    • Allow users to define all required input features first
    • Prevent errors from missing DataFrame columns
    • Give users full control over when inference happens
    • Support complex analysis workflows with dependencies
  6. Multiple Outputs: Full support for models with multiple output tensors to:
    • Enable ParticleTransformer-style models with bootstrapped outputs
    • Support ensemble models that return multiple predictions
    • Create individual DataFrame columns for each output for easy access

Security Considerations

Future Enhancements

Potential improvements for future work:

  1. Support for multiple output tensors from models - Implemented ✅
  2. Support for variable-shaped inputs (dynamic batch sizes, ragged arrays)
  3. Batch inference for improved performance
  4. Model caching for faster repeated evaluations
  5. Support for different ONNX Runtime execution providers (CPU, CUDA, TensorRT)
  6. Model quantization support
  7. Integration with ONNX Model Zoo

Compatibility

Limitations

  1. Float Input: Input features are converted to float32
  2. Fixed Shape Input: Currently expects (1, N) shaped input tensors for single input models
  3. CPU Only: ONNX Runtime is configured for CPU inference only
  4. Single Scalar Output Per Tensor: Each output tensor must produce a single scalar value

References