Class Acts::OnnxEdgeClassifier

class OnnxEdgeClassifier : public Acts::EdgeClassificationBase

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, int deviceHint = -1) override

Perform edge classification.

Parameters

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

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

• deviceHint – Which GPU to pick. Not relevant for CPU-only builds

Returns

(node_tensor, edge_tensor, score_tensor)

struct Config

Public Members

float cut = 0.21

std::string modelPath