From 850dd4373a632b40c009401dab871186014229eb Mon Sep 17 00:00:00 2001 From: Wouter Deconinck Date: Tue, 16 Apr 2024 22:47:06 -0500 Subject: [PATCH] feat: example of how to use ONNX Runtime (Ort) in algorithms (#1358) ### Briefly, what does this PR introduce? This PR adds the hooks for ONNX Runtime (CPU only) for fast inference, and an example that doesn't actually do anything except... Enabled by default but not actually doing anything. ### What kind of change does this PR introduce? - [ ] Bug fix (issue #__) - [x] New feature (issue: use ONNX for ML) - [ ] Documentation update - [ ] Other: __ ### Please check if this PR fulfills the following: - [ ] Tests for the changes have been added - [ ] Documentation has been added / updated - [x] Changes have been communicated to collaborators @rahmans1 ### Does this PR introduce breaking changes? What changes might users need to make to their code? No. ### Does this PR change default behavior? No. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- CMakeLists.txt | 6 + cmake/jana_plugin.cmake | 12 ++ src/algorithms/CMakeLists.txt | 4 + src/algorithms/onnx/CMakeLists.txt | 21 +++ src/algorithms/onnx/InclusiveKinematicsML.cc | 127 ++++++++++++++++++ src/algorithms/onnx/InclusiveKinematicsML.h | 60 +++++++++ .../onnx/InclusiveKinematicsMLConfig.h | 16 +++ .../reco/InclusiveKinematicsML_factory.h | 47 +++++++ src/global/reco/CMakeLists.txt | 5 +- src/global/reco/reco.cc | 14 ++ src/services/io/podio/JEventProcessorPODIO.cc | 1 + 11 files changed, 311 insertions(+), 2 deletions(-) create mode 100644 src/algorithms/onnx/CMakeLists.txt create mode 100644 src/algorithms/onnx/InclusiveKinematicsML.cc create mode 100644 src/algorithms/onnx/InclusiveKinematicsML.h create mode 100644 src/algorithms/onnx/InclusiveKinematicsMLConfig.h create mode 100644 src/factories/reco/InclusiveKinematicsML_factory.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 724b798850..aec87fb4d1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -219,6 +219,12 @@ if(NOT "cxx_std_${CMAKE_CXX_STANDARD}" IN_LIST ROOT_COMPILE_FEATURES) ) endif() +# ONNX Runtime +option(USE_ONNX "Compile with ONNX support" ON) +if(${USE_ONNX}) + find_package(onnxruntime) +endif() + # Add CMake additional functionality: include(cmake/jana_plugin.cmake) # Add common settings for plugins list(APPEND CMAKE_MODULE_PATH ${EICRECON_SOURCE_DIR}/cmake diff --git a/cmake/jana_plugin.cmake b/cmake/jana_plugin.cmake index 8450bcc055..da4454cbb0 100644 --- a/cmake/jana_plugin.cmake +++ b/cmake/jana_plugin.cmake @@ -376,3 +376,15 @@ macro(plugin_add_fastjet _name) plugin_link_libraries(${PLUGIN_NAME} ${FASTJET_LIBRARIES}) endmacro() + +# Adds ONNX Runtime for a plugin +macro(plugin_add_onnxruntime _name) + + if(NOT onnxruntime_FOUND) + find_package(onnxruntime) + endif() + + # Add libraries + plugin_link_libraries(${PLUGIN_NAME} onnxruntime::onnxruntime) + +endmacro() diff --git a/src/algorithms/CMakeLists.txt b/src/algorithms/CMakeLists.txt index dddced5bff..013ea11931 100644 --- a/src/algorithms/CMakeLists.txt +++ b/src/algorithms/CMakeLists.txt @@ -7,3 +7,7 @@ add_subdirectory(pid) add_subdirectory(digi) add_subdirectory(reco) add_subdirectory(fardetectors) + +if(USE_ONNX) + add_subdirectory(onnx) +endif() diff --git a/src/algorithms/onnx/CMakeLists.txt b/src/algorithms/onnx/CMakeLists.txt new file mode 100644 index 0000000000..c0fd746a91 --- /dev/null +++ b/src/algorithms/onnx/CMakeLists.txt @@ -0,0 +1,21 @@ +cmake_minimum_required(VERSION 3.16) + +set(PLUGIN_NAME "algorithms_onnx") + +# Function creates ${PLUGIN_NAME}_plugin and ${PLUGIN_NAME}_library targets +# Setting default includes, libraries and installation paths +plugin_add(${PLUGIN_NAME} WITH_SHARED_LIBRARY WITHOUT_PLUGIN) + +# The macro grabs sources as *.cc *.cpp *.c and headers as *.h *.hh *.hpp Then +# correctly sets sources for ${_name}_plugin and ${_name}_library targets Adds +# headers to the correct installation directory +plugin_glob_all(${PLUGIN_NAME}) + +# Find dependencies +plugin_add_event_model(${PLUGIN_NAME}) +plugin_add_onnxruntime(${PLUGIN_NAME}) + +# The macro grabs sources as *.cc *.cpp *.c and headers as *.h *.hh *.hpp Then +# correctly sets sources for ${_name}_plugin and ${_name}_library targets Adds +# headers to the correct installation directory +plugin_glob_all(${PLUGIN_NAME}) diff --git a/src/algorithms/onnx/InclusiveKinematicsML.cc b/src/algorithms/onnx/InclusiveKinematicsML.cc new file mode 100644 index 0000000000..6ef2e72a40 --- /dev/null +++ b/src/algorithms/onnx/InclusiveKinematicsML.cc @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2022, 2023 Wouter Deconinck, Tooba Ali + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "InclusiveKinematicsML.h" + +namespace eicrecon { + + static std::string print_shape(const std::vector& v) { + std::stringstream ss(""); + for (std::size_t i = 0; i < v.size() - 1; i++) ss << v[i] << "x"; + ss << v[v.size() - 1]; + return ss.str(); + } + + template + Ort::Value vec_to_tensor(std::vector& data, const std::vector& shape) { + Ort::MemoryInfo mem_info = + Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); + auto tensor = Ort::Value::CreateTensor(mem_info, data.data(), data.size(), shape.data(), shape.size()); + return tensor; + } + + void InclusiveKinematicsML::init(std::shared_ptr& logger) { + m_log = logger; + + // onnxruntime setup + Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "inclusive-kinematics-ml"); + Ort::SessionOptions session_options; + try { + m_session = Ort::Session(env, m_cfg.modelPath.c_str(), session_options); + + // print name/shape of inputs + Ort::AllocatorWithDefaultOptions allocator; + m_log->debug("Input Node Name/Shape:"); + for (std::size_t i = 0; i < m_session.GetInputCount(); i++) { + m_input_names.emplace_back(m_session.GetInputNameAllocated(i, allocator).get()); + m_input_shapes.emplace_back(m_session.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape()); + m_log->debug("\t{} : {}", m_input_names.at(i), print_shape(m_input_shapes.at(i))); + } + + // print name/shape of outputs + m_log->debug("Output Node Name/Shape:"); + for (std::size_t i = 0; i < m_session.GetOutputCount(); i++) { + m_output_names.emplace_back(m_session.GetOutputNameAllocated(i, allocator).get()); + m_output_shapes.emplace_back(m_session.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape()); + m_log->debug("\t{} : {}", m_output_names.at(i), print_shape(m_output_shapes.at(i))); + } + + // convert names to char* + m_input_names_char.resize(m_input_names.size(), nullptr); + std::transform(std::begin(m_input_names), std::end(m_input_names), std::begin(m_input_names_char), + [&](const std::string& str) { return str.c_str(); }); + m_output_names_char.resize(m_output_names.size(), nullptr); + std::transform(std::begin(m_output_names), std::end(m_output_names), std::begin(m_output_names_char), + [&](const std::string& str) { return str.c_str(); }); + + } catch(std::exception& e) { + m_log->error(e.what()); + } + } + + void InclusiveKinematicsML::process( + const InclusiveKinematicsML::Input& input, + const InclusiveKinematicsML::Output& output) const { + + const auto [electron, da] = input; + auto [ml] = output; + + // Require valid inputs + if (electron->size() == 0 || da->size() == 0) { + m_log->debug("skipping because input collections have no entries"); + return; + } + + // Assume model has 1 input nodes and 1 output node. + if (m_input_names.size() != 1 || m_output_names.size() != 1) { + m_log->debug("skipping because model has incorrect input and output size"); + return; + } + + // Prepare input tensor + std::vector input_tensor_values; + std::vector input_tensors; + for (std::size_t i = 0; i < electron->size(); i++) { + input_tensor_values.push_back(electron->at(i).getX()); + } + input_tensors.emplace_back(vec_to_tensor(input_tensor_values, m_input_shapes.front())); + + // Double-check the dimensions of the input tensor + if (! input_tensors[0].IsTensor() || input_tensors[0].GetTensorTypeAndShapeInfo().GetShape() != m_input_shapes.front()) { + m_log->debug("skipping because input tensor shape incorrect"); + return; + } + + // Attempt inference + try { + auto output_tensors = m_session.Run(Ort::RunOptions{nullptr}, m_input_names_char.data(), input_tensors.data(), + m_input_names_char.size(), m_output_names_char.data(), m_output_names_char.size()); + + // Double-check the dimensions of the output tensors + if (!output_tensors[0].IsTensor() || output_tensors.size() != m_output_names.size()) { + m_log->debug("skipping because output tensor shape incorrect"); + return; + } + + // Convert output tensor + float* output_tensor_data = output_tensors[0].GetTensorMutableData(); + auto x = output_tensor_data[0]; + auto kin = ml->create(); + kin.setX(x); + + } catch (const Ort::Exception& exception) { + m_log->error("error running model inference: {}", exception.what()); + } + } + +} // namespace eicrecon diff --git a/src/algorithms/onnx/InclusiveKinematicsML.h b/src/algorithms/onnx/InclusiveKinematicsML.h new file mode 100644 index 0000000000..e2d559dbf0 --- /dev/null +++ b/src/algorithms/onnx/InclusiveKinematicsML.h @@ -0,0 +1,60 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2022, 2023 Sylvester Joosten, Dmitry Romanov, Wouter Deconinck + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "algorithms/interfaces/WithPodConfig.h" +#include "algorithms/onnx/InclusiveKinematicsMLConfig.h" + +namespace eicrecon { + + using InclusiveKinematicsMLAlgorithm = algorithms::Algorithm< + algorithms::Input< + edm4eic::InclusiveKinematicsCollection, + edm4eic::InclusiveKinematicsCollection + >, + algorithms::Output< + edm4eic::InclusiveKinematicsCollection + > + >; + + class InclusiveKinematicsML + : public InclusiveKinematicsMLAlgorithm, + public WithPodConfig { + + public: + InclusiveKinematicsML(std::string_view name) + : InclusiveKinematicsMLAlgorithm{name, + {"inclusiveKinematicsElectron", "inclusiveKinematicsDA"}, + {"inclusiveKinematicsML"}, + "Determine inclusive kinematics using combined ML method."} {} + + void init(std::shared_ptr& logger); + void process(const Input&, const Output&) const final; + + private: + std::shared_ptr m_log; + + mutable Ort::Session m_session{nullptr}; + + std::vector m_input_names; + std::vector m_input_names_char; + std::vector> m_input_shapes; + + std::vector m_output_names; + std::vector m_output_names_char; + std::vector> m_output_shapes; + + }; + +} // namespace eicrecon diff --git a/src/algorithms/onnx/InclusiveKinematicsMLConfig.h b/src/algorithms/onnx/InclusiveKinematicsMLConfig.h new file mode 100644 index 0000000000..dd7bd37d1a --- /dev/null +++ b/src/algorithms/onnx/InclusiveKinematicsMLConfig.h @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024 Wouter Deconinck + +#pragma once + +#include + +namespace eicrecon { + + struct InclusiveKinematicsMLConfig { + + std::string modelPath{"calibrations/onnx/identity_gemm_w1x1_b1.onnx"}; + + }; + +} // eicrecon diff --git a/src/factories/reco/InclusiveKinematicsML_factory.h b/src/factories/reco/InclusiveKinematicsML_factory.h new file mode 100644 index 0000000000..ad9861289d --- /dev/null +++ b/src/factories/reco/InclusiveKinematicsML_factory.h @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2022 Wouter Deconinck + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "algorithms/onnx/InclusiveKinematicsML.h" +#include "extensions/jana/JOmniFactory.h" + +namespace eicrecon { + +class InclusiveKinematicsML_factory : + public JOmniFactory { + +public: + using AlgoT = eicrecon::InclusiveKinematicsML; +private: + std::unique_ptr m_algo; + + PodioInput m_inclusive_kinematics_electron_input {this}; + PodioInput m_inclusive_kinematics_da_input {this}; + PodioOutput m_inclusive_kinematics_output {this}; + + ParameterRef m_modelPath {this, "modelPath", config().modelPath}; + +public: + void Configure() { + m_algo = std::make_unique(GetPrefix()); + m_algo->applyConfig(config()); + m_algo->init(logger()); + } + + void ChangeRun(int64_t run_number) { + } + + void Process(int64_t run_number, uint64_t event_number) { + m_algo->process({m_inclusive_kinematics_electron_input(), m_inclusive_kinematics_da_input()}, {m_inclusive_kinematics_output().get()}); + } +}; + +} // eicrecon diff --git a/src/global/reco/CMakeLists.txt b/src/global/reco/CMakeLists.txt index 5049d1ac7a..90886e8147 100644 --- a/src/global/reco/CMakeLists.txt +++ b/src/global/reco/CMakeLists.txt @@ -26,5 +26,6 @@ plugin_glob_all(${PLUGIN_NAME}) # Add libraries (same as target_include_directories but for both plugin and # library) -plugin_link_libraries(${PLUGIN_NAME} algorithms_digi_library - algorithms_tracking_library algorithms_reco_library) +plugin_link_libraries( + ${PLUGIN_NAME} algorithms_digi_library algorithms_tracking_library + algorithms_onnx_library algorithms_reco_library) diff --git a/src/global/reco/reco.cc b/src/global/reco/reco.cc index 5a3c27b59c..fd13ec7b38 100644 --- a/src/global/reco/reco.cc +++ b/src/global/reco/reco.cc @@ -11,6 +11,7 @@ #include #include +#include "algorithms/interfaces/WithPodConfig.h" #include "algorithms/reco/InclusiveKinematicsDA.h" #include "algorithms/reco/InclusiveKinematicsElectron.h" #include "algorithms/reco/InclusiveKinematicsJB.h" @@ -18,6 +19,7 @@ #include "algorithms/reco/InclusiveKinematicseSigma.h" #include "extensions/jana/JOmniFactoryGeneratorT.h" #include "factories/meta/CollectionCollector_factory.h" +#include "factories/reco/InclusiveKinematicsML_factory.h" #include "factories/reco/InclusiveKinematicsReconstructed_factory.h" #include "factories/reco/InclusiveKinematicsTruth_factory.h" #include "factories/reco/JetReconstruction_factory.h" @@ -155,6 +157,18 @@ void InitPlugin(JApplication *app) { app )); + app->Add(new JOmniFactoryGeneratorT( + "InclusiveKinematicsML", + { + "InclusiveKinematicsElectron", + "InclusiveKinematicsDA" + }, + { + "InclusiveKinematicsML" + }, + app + )); + app->Add(new JOmniFactoryGeneratorT( "ReconstructedElectrons", {"ReconstructedParticles"}, diff --git a/src/services/io/podio/JEventProcessorPODIO.cc b/src/services/io/podio/JEventProcessorPODIO.cc index 8fca53c557..698b97f990 100644 --- a/src/services/io/podio/JEventProcessorPODIO.cc +++ b/src/services/io/podio/JEventProcessorPODIO.cc @@ -133,6 +133,7 @@ JEventProcessorPODIO::JEventProcessorPODIO() { "CentralCKFSeededTrackParameters", "InclusiveKinematicsDA", "InclusiveKinematicsJB", + "InclusiveKinematicsML", "InclusiveKinematicsSigma", "InclusiveKinematicseSigma", "InclusiveKinematicsElectron",