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",