File OnnxEdgeClassifier.hpp

namespace Acts

Note

This file is foreseen for the Geometry module to replace Extent

class OnnxEdgeClassifier : public Acts::EdgeClassificationBase

#include <Acts/Plugins/ExaTrkX/OnnxEdgeClassifier.hpp>

Public Functions

OnnxEdgeClassifier(const Config &cfg, std::unique_ptr<const Logger> logger)

~OnnxEdgeClassifier()

inline Config config() const

virtual std::tuple<std::any, std::any, std::any> operator()(std::any nodes, std::any edges) override

Perform edge classification.

Parameters

• nodes – Node tensor with shape (n_nodes, n_node_features)

• edges – Edge-index tensor with shape (2, n_edges)

Returns

(node_tensor, edge_tensor, score_tensor)

Private Functions

inline const auto &logger() const

Private Members

Config m_cfg

std::unique_ptr<Ort::Env> m_env

std::string m_inputNameEdges

std::string m_inputNameNodes

std::unique_ptr<const Acts::Logger> m_logger

std::unique_ptr<Ort::Session> m_model

std::string m_outputNameScores

struct Config

#include <Acts/Plugins/ExaTrkX/OnnxEdgeClassifier.hpp>

Public Members

float cut = 0.21

std::string modelPath

namespace Ort