Plugins/Gnn/OnnxEdgeClassifierTests.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #include <boost/test/unit_test.hpp>
0010
0011 #include "Acts/Utilities/Logger.hpp"
0012 #include "ActsPlugins/Gnn/OnnxEdgeClassifier.hpp"
0013 #include "ActsTests/CommonHelpers/DataDirectory.hpp"
0014
0015 #include <algorithm>
0016 #include <cmath>
0017 #include <cstdint>
0018 #include <stdexcept>
0019 #include <string_view>
0020 #include <vector>
0021
0022 #ifdef ACTS_GNN_WITH_CUDA
0023 #include <cuda_runtime_api.h>
0024 #endif
0025
0026 using namespace ActsPlugins;
0027
0028 namespace ActsTests {
0029
0030 namespace {
0031
0032 const ExecutionContext execContextCpu{Device::Cpu(), {}};
0033 #ifdef ACTS_GNN_WITH_CUDA
0034 const ExecutionContext execContextCuda{Device::Cuda(0), cudaStreamLegacy};
0035 #endif
0036
0037 // The classifier runs echo_edge_classifier.onnx (see
0038 // make_echo_edge_classifier.py), which returns the node features it is given,
0039 // flattened row by row, as its edge scores. Undoing the sigmoid the classifier
0040 // applies to them gives back the feature matrix that reached the model.
0041
0042 // Two nodes with four features each, e.g. (r, phi, z, t). All values differ,
0043 // so any swapped column or misapplied scale shows up.
0044 constexpr std::size_t nNodes = 2;
0045 constexpr std::size_t nFeatures = 4;
0046 const std::vector<float> nodes = {
0047     0.5f, -1.5f, 2.5f,  0.25f,  // node 0
0048     1.0f, 0.75f, -3.0f, 1.75f   // node 1
0049 };
0050
0051 /// Run the echo classifier with @p selectedFeatures and @p featureScales and
0052 /// return the (row-major) node feature matrix the model got to see
0053 std::vector<float> modelInputs(const std::vector<int> &selectedFeatures,
0054                                const std::vector<float> &featureScales,
0055                                const ExecutionContext &execContext) {
0056   OnnxEdgeClassifier::Config cfg;
0057   cfg.modelPath = getDataPath("echo_edge_classifier.onnx");
0058   cfg.selectedFeatures = selectedFeatures;
0059   cfg.featureScales = featureScales;
0060   // Keep every edge, the sigmoid is always above
0061   cfg.cut = 0.f;
0062   cfg.device = execContext.device;
0063   OnnxEdgeClassifier classifier(
0064       cfg,
0065       Acts::getDefaultLogger("OnnxEdgeClassifier", Acts::Logging::WARNING));
0066
0067   auto nodeFeatures =
0068       Tensor<float>::Create({nNodes, nFeatures}, execContextCpu);
0069   std::ranges::copy(nodes, nodeFeatures.data());
0070
0071   // One edge per value the model reads. The edge index itself only sets the
0072   // number of scores the model returns.
0073   const std::size_t nModelFeatures =
0074       selectedFeatures.empty() ? nFeatures : selectedFeatures.size();
0075   const std::size_t nEdges = nNodes * nModelFeatures;
0076   auto edgeIndex = Tensor<std::int64_t>::Create({2, nEdges}, execContextCpu);
0077   std::fill_n(edgeIndex.data(), edgeIndex.size(), 0);
0078
0079   PipelineTensors tensors{
0080       nodeFeatures.clone(execContext), edgeIndex.clone(execContext), {}, {}};
0081   auto out = classifier(std::move(tensors), execContext);
0082
0083   // Copying back from the device needs the stream, a CPU context has none
0084   const ExecutionContext toHost{Device::Cpu(), execContext.stream};
0085
0086   // The stage passes the full, unscaled node features on to the next one
0087   auto outNodes = out.nodeFeatures.clone(toHost);
0088   BOOST_CHECK_EQUAL_COLLECTIONS(outNodes.data(),
0089                                 outNodes.data() + outNodes.size(),
0090                                 nodes.begin(), nodes.end());
0091
0092   BOOST_REQUIRE(out.edgeScores.has_value());
0093   auto scores = out.edgeScores->clone(toHost);
0094   BOOST_REQUIRE_EQUAL(scores.size(), nEdges);
0095   std::vector<float> inputs(nEdges);
0096   std::transform(scores.data(), scores.data() + nEdges, inputs.begin(),
0097                  [](float s) { return std::log(s / (1.f - s)); });
0098   return inputs;
0099 }
0100
0101 void checkInputs(const std::vector<float> &actual,
0102                  const std::vector<float> &expected) {
0103   BOOST_REQUIRE_EQUAL(actual.size(), expected.size());
0104   for (std::size_t i = 0; i < actual.size(); ++i) {
0105     BOOST_TEST_CONTEXT("model input " << i) {
0106       BOOST_CHECK_SMALL(actual[i] - expected[i], 1e-4f);
0107     }
0108   }
0109 }
0110
0111 void testFeatureSelection(const ExecutionContext &ctx) {
0112   // Without a selection the model gets all features
0113   checkInputs(modelInputs({}, {}, ctx), nodes);
0114
0115   // The features reach the model in the configured order: (z, r), not (r, z)
0116   // as they are ordered in the node features
0117   checkInputs(modelInputs({2, 0}, {}, ctx), {2.5f, 0.5f, -3.f, 1.f});
0118
0119   // Each scale divides its own feature: z / 10 and r / 2
0120   checkInputs(modelInputs({2, 0}, {10.f, 2.f}, ctx),
0121               {0.25f, 0.25f, -0.3f, 0.5f});
0122
0123   // Without a selection the scales apply to the features in order
0124   checkInputs(modelInputs({}, {1.f, 2.f, 4.f, 0.5f}, ctx),
0125               {0.5f, -0.75f, 0.625f, 0.5f, 1.f, 0.375f, -0.75f, 3.5f});
0126
0127   // A feature can be selected more than once: phi twice, with a different
0128   // scale each time, around t
0129   checkInputs(modelInputs({1, 3, 1}, {1.f, 1.f, 3.f}, ctx),
0130               {-1.5f, 0.25f, -0.5f, 0.75f, 1.75f, 0.25f});
0131
0132   // A selected feature has to exist
0133   BOOST_CHECK_THROW(modelInputs({0, 4}, {}, ctx), std::runtime_error);
0134   BOOST_CHECK_THROW(modelInputs({-1}, {}, ctx), std::runtime_error);
0135
0136   // There has to be one scale per model input feature
0137   BOOST_CHECK_THROW(modelInputs({2, 0}, {1.f}, ctx), std::runtime_error);
0138   BOOST_CHECK_THROW(modelInputs({}, {1.f}, ctx), std::runtime_error);
0139 }
0140
0141 }  // namespace
0142
0143 BOOST_AUTO_TEST_SUITE(GnnOnnxEdgeClassifierSuite)
0144
0145 BOOST_AUTO_TEST_CASE(test_feature_selection_cpu) {
0146   testFeatureSelection(execContextCpu);
0147 }
0148
0149 #ifdef ACTS_GNN_WITH_CUDA
0150 BOOST_AUTO_TEST_CASE(test_feature_selection_cuda) {
0151   testFeatureSelection(execContextCuda);
0152 }
0153 #endif
0154
0155 BOOST_AUTO_TEST_CASE(test_missing_model_is_named) {
0156   OnnxEdgeClassifier::Config cfg;
0157   cfg.modelPath = "this-model-does-not-exist.onnx";
0158   cfg.device = Device::Cpu();
0159   BOOST_CHECK_EXCEPTION(
0160       OnnxEdgeClassifier(cfg, Acts::getDefaultLogger("OnnxEdgeClassifier",
0161                                                      Acts::Logging::WARNING)),
0162       std::runtime_error, [&cfg](const std::runtime_error &e) {
0163         return std::string_view{e.what()}.find(cfg.modelPath) !=
0164                std::string_view::npos;
0165       });
0166 }
0167
0168 BOOST_AUTO_TEST_SUITE_END()
0169
0170 }  // namespace ActsTests