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. More...
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 More...

Private Attributes
NodeFeatureList	m_graphFeatures {}
	List of features to be used for the inference. More...
ToolHandle< AthOnnx::IOnnxRuntimeSessionTool >	m_onnxSessionTool {this, "ModelSession", "" }

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 114 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 87 of file SPInferenceToolBase.cxx.

                                                 {
        return m_onnxSessionTool->session();
    }

runInference()

StatusCode MuonML::SPInferenceToolBase::runInference

(

GraphRawData &

graphData

)

const

Definition at line 188 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_WARNING);
 
        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 90 of file SPInferenceToolBase.cxx.

                                               {
        ATH_CHECK(m_onnxSessionTool.retrieve());
        ATH_CHECK(m_readKey.initialize());
            
        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_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_onnxSessionTool

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

private

Definition at line 45 of file SPInferenceToolBase.h.

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.

---

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

• SPInferenceToolBase.h
• SPInferenceToolBase.cxx