MuonML::SPInferenceToolBase Class Reference

Baseline tool to handle the. More...

#include <SPInferenceToolBase.h>

Inheritance diagram for MuonML::SPInferenceToolBase:

Collaboration diagram for MuonML::SPInferenceToolBase:

Public Member Functions
StatusCode	buildGraph (const EventContext &ctx, GraphRawData &graphData) const
	Fill up the GraphRawData and construct the graph for the ML inference with ONNX.
StatusCode	runInference (GraphRawData &graphData) const

Protected Member Functions
StatusCode	setupModel ()
Ort::Session &	model () const

Protected Attributes
SG::ReadHandleKey< MuonR4::SpacePointContainer >	m_readKey {this, "ReadSpacePoints", "MuonSpacePoints"}
	Input space points to filter.

Private Attributes
NodeFeatureList	m_graphFeatures {}
	List of features to be used for the inference.
ToolHandle< AthOnnx::IOnnxRuntimeSessionTool >	m_onnxSessionTool {this, "ModelSession", "" }
bool	m_isCuda {false}
int	m_cudaDeviceId {0}

Detailed Description

Baseline tool to handle the.

Definition at line 20 of file SPInferenceToolBase.h.

Member Function Documentation

buildGraph()

StatusCode MuonML::SPInferenceToolBase::buildGraph	(	const EventContext &	ctx,
		GraphRawData &	graphData ) const

Fill up the GraphRawData and construct the graph for the ML inference with ONNX.

If the graph has been built by another inference tool and would be the same than this one the rebuild is skipped

Parameters

ctx	EventContext to access the space ponit container from StoreGate
graphData	Rerference to the data object to be filled.

Check whether the graph needs a rebuild

Don't launch the rebuild of the graph

Fill the graph edge features and all their respective connections

Definition at line 127 of file SPInferenceToolBase.cxx.

                                                                                   {

        if (graphData.previousList && (*graphData.previousList) != m_graphFeatures) {
            graphData.graph.reset();
        }
        if (graphData.graph) {
            return StatusCode::SUCCESS;
        }        
        if (!m_graphFeatures.isValid()) {
            ATH_MSG_ERROR("The feature list is in complete. Either it has no features or no node connector set");
            return StatusCode::FAILURE;
        }
    
        graphData.graph = std::make_unique<InferenceGraph>();
    
        SG::ReadHandle spacePoints{m_readKey, ctx};
        ATH_CHECK(spacePoints.isPresent());
    
        int64_t nNodes{0}, possConn{0};
        graphData.spacePointsInBucket.clear();
        graphData.spacePointsInBucket.reserve(spacePoints->size());
 
        for (const MuonR4::SpacePointBucket* bucket : *spacePoints) {
            nNodes += graphData.spacePointsInBucket.emplace_back(bucket->size());
            possConn += Acts::sumUpToN(graphData.spacePointsInBucket.back());
        }
 
        graphData.nodeIndex = 0;
        graphData.featureLeaves.resize(nNodes * m_graphFeatures.numFeatures());
        graphData.currLeave = graphData.featureLeaves.begin();
        
        graphData.srcEdges.reserve(possConn);
        graphData.desEdges.reserve(possConn);

        for (const MuonR4::SpacePointBucket* bucket : *spacePoints) {
            const LayerSpBucket mlBucket{*bucket};
            m_graphFeatures.fillInData(mlBucket, graphData);
        }
 
        graphData.srcEdges.insert(graphData.srcEdges.end(), std::make_move_iterator(graphData.desEdges.begin()),
                                    std::make_move_iterator(graphData.desEdges.end()));
 
        ATH_MSG_DEBUG("Features:"<<m_graphFeatures.featureNames());
        ATH_MSG_DEBUG(formatNodeFeatures(graphData.featureLeaves, m_graphFeatures.numFeatures()));
        ATH_MSG_DEBUG("Edge indices:"<<makeIndexPairs(makeSortedEdges(graphData.srcEdges)));
 
        std::vector<int64_t> featShape{nNodes, static_cast<int64_t>(m_graphFeatures.numFeatures())};    // (N, nFeatures)
        std::vector<int64_t> edgeShape{2, static_cast<int64_t>(graphData.srcEdges.size() / 2)};         // (2, E)
 
        
        Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
        graphData.graph->dataTensor.emplace_back(Ort::Value::CreateTensor<float>(memInfo, 
                                                    graphData.featureLeaves.data(), graphData.featureLeaves.size(), 
                                                    featShape.data(), featShape.size()));
 
 
        Ort::Value edge_tensor = Ort::Value::CreateTensor<int64_t>(memInfo, graphData.srcEdges.data(), 
            graphData.srcEdges.size(), edgeShape.data(), edgeShape.size());
 
        graphData.graph->dataTensor.emplace_back(std::move(edge_tensor));
 
        graphData.previousList = &m_graphFeatures;
 
        graphData.srcEdges.clear();
        graphData.desEdges.clear();
        graphData.featureLeaves.clear();
        ATH_MSG_DEBUG("Graph data built successfully.");
        return StatusCode::SUCCESS;
    }

model()

Ort::Session & MuonML::SPInferenceToolBase::model ( ) const

protected

Definition at line 88 of file SPInferenceToolBase.cxx.

                                                 {
        return m_onnxSessionTool->session();
    }

runInference()

StatusCode MuonML::SPInferenceToolBase::runInference

(

GraphRawData &

graphData

)

const

Definition at line 201 of file SPInferenceToolBase.cxx.

                                                                              {
        if (!m_graphFeatures.isValid()) {
            ATH_MSG_ERROR("ONNX model is not loaded. Please call setupModel()");
            return StatusCode::FAILURE;
        }
        if (!graphData.graph) {
            ATH_MSG_ERROR("Graph data is not built.");
            return StatusCode::FAILURE;
        }
 
        if (graphData.graph->dataTensor.size() < 2) {
            ATH_MSG_ERROR("Data tensor does not contain both feature and edge tensors.");
            return StatusCode::FAILURE;
        }
 
        std::vector<const char*> inputNames  = {"features", "edge_index"};
        std::vector<const char*> outputNames = {"output"};
 
        Ort::RunOptions run_options;
        run_options.SetRunLogSeverityLevel(ORT_LOGGING_LEVEL_ERROR);
 
        if (m_isCuda) {
            // ---- CUDA path: use IoBinding ----
            Ort::IoBinding binding(model());
            for (std::size_t i = 0; i < inputNames.size(); ++i) {
                binding.BindInput(inputNames[i], graphData.graph->dataTensor[i]);
            }
            // Bind output to CPU so logits are directly readable after sync.
            Ort::MemoryInfo cpuOut = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
            for (const char* outName : outputNames) {
                binding.BindOutput(outName, cpuOut);
            }
 
            model().Run(run_options, binding);
            binding.SynchronizeOutputs();
 
            std::vector<Ort::Value> outputTensors = binding.GetOutputValues();
            if (outputTensors.empty()) {
                ATH_MSG_ERROR("IoBinding inference returned empty output.");
                return StatusCode::FAILURE;
            }
 
            float* output_data = outputTensors[0].GetTensorMutableData<float>();
            size_t output_size = outputTensors[0].GetTensorTypeAndShapeInfo().GetElementCount();
 
            std::span<float> predictions(output_data, output_data + output_size);
            for (size_t i = 0; i < output_size; i++) {
                if (!std::isfinite(predictions[i])) {
                    ATH_MSG_WARNING("Non-finite prediction detected! Setting to -100..");
                    predictions[i] = -100.0f;
                }
            }
            graphData.graph->dataTensor.emplace_back(std::move(outputTensors[0]));
            return StatusCode::SUCCESS;
        }
 
        // ---- CPU path (unchanged) ----
        std::vector<Ort::Value> outputTensors = model().Run(run_options, 
                                                inputNames.data(),                    // input tensor names
                                                graphData.graph->dataTensor.data(),  // pointer to the tensor vector
                                                graphData.graph->dataTensor.size(),   // size of the tensor vector
                                                outputNames.data(),                  // output tensor names
                                                outputNames.size());                 // number of output tensors
    
        if (outputTensors.empty()) {
            ATH_MSG_ERROR("Inference returned empty output.");
            return StatusCode::FAILURE;
        }     
 
        float* output_data = outputTensors[0].GetTensorMutableData<float>();
        size_t output_size = outputTensors[0].GetTensorTypeAndShapeInfo().GetElementCount();
 
        std::span<float> predictions(output_data, output_data + output_size);
        
        for (size_t i = 0; i < output_size; i++) {
            if (!std::isfinite(predictions[i])) {
                ATH_MSG_WARNING("Non-finite prediction detected! Setting to -100..");
                predictions[i] = -100.0f;
            }
        }
 
        graphData.graph->dataTensor.emplace_back(std::move(outputTensors[0]));
 
        return StatusCode::SUCCESS;
    }

setupModel()

StatusCode MuonML::SPInferenceToolBase::setupModel ( )

protected

Definition at line 91 of file SPInferenceToolBase.cxx.

                                               {
        ATH_CHECK(m_onnxSessionTool.retrieve());
        ATH_CHECK(m_readKey.initialize());
 
        // Detect CUDA provider by dynamic-casting the concrete session tool.
        if (const auto* cudaTool = dynamic_cast<const AthOnnx::OnnxRuntimeSessionToolCUDA*>(
                m_onnxSessionTool.get())) {
            m_isCuda       = true;
            m_cudaDeviceId = cudaTool->deviceId();
            ATH_MSG_DEBUG("ONNX session is running on CUDA device " << m_cudaDeviceId
                         << ". I/O binding will be used.");
        } else {
            m_isCuda = false;
            ATH_MSG_DEBUG("ONNX session is running on CPU.");
        }
 
        Ort::ModelMetadata metadata = model().GetModelMetadata();
        Ort::AllocatorWithDefaultOptions allocator;
        Ort::AllocatedStringPtr feature_json_ptr = metadata.LookupCustomMetadataMapAllocated("feature_names", allocator);
            
        if (feature_json_ptr) {
            std::string feature_json = feature_json_ptr.get();
            nlohmann::json json_obj = nlohmann::json::parse(feature_json);
            for (const auto& feature : json_obj) {
                m_graphFeatures.addFeature(feature.get<std::string>(), msgStream());
            }
        }
        m_graphFeatures.setConnector("fullyConnected", msgStream());
 
        if (!m_graphFeatures.isValid()) {
            ATH_MSG_FATAL("No graph features have been parsed. Please check the model: "<<m_graphFeatures.featureNames());
            return StatusCode::FAILURE;
        }
        return StatusCode::SUCCESS;
    }

Member Data Documentation

m_cudaDeviceId

int MuonML::SPInferenceToolBase::m_cudaDeviceId {0}

private

Definition at line 48 of file SPInferenceToolBase.h.

{0};

m_graphFeatures

NodeFeatureList MuonML::SPInferenceToolBase::m_graphFeatures {}

private

List of features to be used for the inference.

Definition at line 44 of file SPInferenceToolBase.h.

{};

m_isCuda

bool MuonML::SPInferenceToolBase::m_isCuda {false}

private

Definition at line 47 of file SPInferenceToolBase.h.

{false};

m_onnxSessionTool

ToolHandle<AthOnnx::IOnnxRuntimeSessionTool> MuonML::SPInferenceToolBase::m_onnxSessionTool {this, "ModelSession", "" }

private

Definition at line 45 of file SPInferenceToolBase.h.

{this, "ModelSession", "" };

m_readKey

SG::ReadHandleKey<MuonR4::SpacePointContainer> MuonML::SPInferenceToolBase::m_readKey {this, "ReadSpacePoints", "MuonSpacePoints"}

protected

Input space points to filter.

Definition at line 41 of file SPInferenceToolBase.h.

{this, "ReadSpacePoints", "MuonSpacePoints"};

---

The documentation for this class was generated from the following files:

• SPInferenceToolBase.h
• SPInferenceToolBase.cxx