From 80f3eb6b285eb9e4680b4e772d7d12df053e78df Mon Sep 17 00:00:00 2001 From: Dmitry Kalinkin Date: Fri, 30 Aug 2024 14:50:02 -0400 Subject: [PATCH 1/2] add ONNXInference algorithm, use it to provide EcalEndcapNClusterParticleIDs --- .../onnx/CalorimeterParticleIDPostML.cc | 86 ++++++++ .../onnx/CalorimeterParticleIDPostML.h | 44 ++++ .../onnx/CalorimeterParticleIDPreML.cc | 100 ++++++++++ .../onnx/CalorimeterParticleIDPreML.h | 39 ++++ src/algorithms/onnx/ONNXInference.cc | 188 ++++++++++++++++++ src/algorithms/onnx/ONNXInference.h | 50 +++++ src/algorithms/onnx/ONNXInferenceConfig.h | 16 ++ src/detectors/EEMC/CMakeLists.txt | 1 + src/detectors/EEMC/EEMC.cc | 59 ++++++ .../CalorimeterParticleIDPostML_factory.h | 45 +++++ .../CalorimeterParticleIDPreML_factory.h | 43 ++++ src/factories/meta/ONNXInference_factory.h | 55 +++++ 12 files changed, 726 insertions(+) create mode 100644 src/algorithms/onnx/CalorimeterParticleIDPostML.cc create mode 100644 src/algorithms/onnx/CalorimeterParticleIDPostML.h create mode 100644 src/algorithms/onnx/CalorimeterParticleIDPreML.cc create mode 100644 src/algorithms/onnx/CalorimeterParticleIDPreML.h create mode 100644 src/algorithms/onnx/ONNXInference.cc create mode 100644 src/algorithms/onnx/ONNXInference.h create mode 100644 src/algorithms/onnx/ONNXInferenceConfig.h create mode 100644 src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h create mode 100644 src/factories/calorimetry/CalorimeterParticleIDPreML_factory.h create mode 100644 src/factories/meta/ONNXInference_factory.h diff --git a/src/algorithms/onnx/CalorimeterParticleIDPostML.cc b/src/algorithms/onnx/CalorimeterParticleIDPostML.cc new file mode 100644 index 0000000000..7c2b2199ab --- /dev/null +++ b/src/algorithms/onnx/CalorimeterParticleIDPostML.cc @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024 Dmitry Kalinkin + +#include + +#if EDM4EIC_VERSION_MAJOR >= 8 +#include +#include +#include +#include + +#include "CalorimeterParticleIDPostML.h" + +namespace eicrecon { + + void CalorimeterParticleIDPostML::init() { + // Nothing + } + + void CalorimeterParticleIDPostML::process( + const CalorimeterParticleIDPostML::Input& input, + const CalorimeterParticleIDPostML::Output& output) const { + + const auto [in_clusters, in_assocs, prediction_tensors] = input; + auto [out_clusters, out_assocs, out_particle_ids] = output; + + if (prediction_tensors->size() != 1) { + error("Expected to find a single tensor, found {}", prediction_tensors->size()); + throw std::runtime_error(""); + } + edm4eic::Tensor prediction_tensor = (*prediction_tensors)[0]; + + if (prediction_tensor.shape_size() != 2) { + error("Expected tensor rank to be 2, but it is {}", prediction_tensor.shape_size()); + throw std::runtime_error(fmt::format("Expected tensor rank to be 2, but it is {}", prediction_tensor.shape_size())); + } + + if (prediction_tensor.getShape(0) != in_clusters->size()) { + error("Length mismatch between tensor's 0th axis and number of clusters: {} != {}", prediction_tensor.getShape(0), in_clusters->size()); + throw std::runtime_error(fmt::format("Length mismatch between tensor's 0th axis and number of clusters: {} != {}", prediction_tensor.getShape(0), in_clusters->size())); + } + + if (prediction_tensor.getShape(1) != 2) { + error("Expected 2 values per cluster in the output tensor, got {}", prediction_tensor.getShape(0)); + throw std::runtime_error(fmt::format("Expected 2 values per cluster in the output tensor, got {}", prediction_tensor.getShape(0))); + } + + if (prediction_tensor.getElementType() != 1) { // 1 - float + error("Expected a tensor of floats, but element type is {}", prediction_tensor.getElementType()); + throw std::runtime_error(fmt::format("Expected a tensor of floats, but element type is {}", prediction_tensor.getElementType())); + } + + for (size_t cluster_ix = 0; cluster_ix < in_clusters->size(); cluster_ix++) { + edm4eic::Cluster in_cluster = (*in_clusters)[cluster_ix]; + edm4eic::MutableCluster out_cluster = in_cluster.clone(); + out_clusters->push_back(out_cluster); + + float prob_pion = prediction_tensor.getFloatData(cluster_ix * prediction_tensor.getShape(1) + 0); + float prob_electron = prediction_tensor.getFloatData(cluster_ix * prediction_tensor.getShape(1) + 1); + + out_cluster.addToParticleIDs(out_particle_ids->create( + 0, // std::int32_t type + 211, // std::int32_t PDG + 0, // std::int32_t algorithmType + prob_pion // float likelihood + )); + out_cluster.addToParticleIDs(out_particle_ids->create( + 0, // std::int32_t type + 11, // std::int32_t PDG + 0, // std::int32_t algorithmType + prob_electron // float likelihood + )); + + // propagate associations + for (auto in_assoc : *in_assocs) { + if (in_assoc.getRec() == in_cluster) { + auto out_assoc = in_assoc.clone(); + out_assoc.setRec(out_cluster); + out_assocs->push_back(out_assoc); + } + } + } + } + +} // namespace eicrecon +#endif diff --git a/src/algorithms/onnx/CalorimeterParticleIDPostML.h b/src/algorithms/onnx/CalorimeterParticleIDPostML.h new file mode 100644 index 0000000000..dbc2cb93e8 --- /dev/null +++ b/src/algorithms/onnx/CalorimeterParticleIDPostML.h @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024 Dmitry Kalinkin + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "algorithms/interfaces/WithPodConfig.h" + +namespace eicrecon { + +using CalorimeterParticleIDPostMLAlgorithm = + algorithms::Algorithm< + algorithms::Input, + edm4eic::TensorCollection>, + algorithms::Output, + edm4hep::ParticleIDCollection> + >; + +class CalorimeterParticleIDPostML : public CalorimeterParticleIDPostMLAlgorithm, + public WithPodConfig { + +public: + CalorimeterParticleIDPostML(std::string_view name) + : CalorimeterParticleIDPostMLAlgorithm{name, + {"inputClusters", "inputClusterAssociations", "inputPredictionsTensor"}, + {"outputClusters", "outputClusterAssociations", "outputParticleIDs"}, + ""} { + } + + void init() final; + void process(const Input&, const Output&) const final; +}; + +} // namespace eicrecon diff --git a/src/algorithms/onnx/CalorimeterParticleIDPreML.cc b/src/algorithms/onnx/CalorimeterParticleIDPreML.cc new file mode 100644 index 0000000000..89d2175d7c --- /dev/null +++ b/src/algorithms/onnx/CalorimeterParticleIDPreML.cc @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024 Dmitry Kalinkin + +#include + +#if EDM4EIC_VERSION_MAJOR >= 8 +#include +#include +#include +#include +#include +#include +#include +#include + +#include "CalorimeterParticleIDPreML.h" + +namespace eicrecon { + + void CalorimeterParticleIDPreML::init() { + // Nothing + } + + void CalorimeterParticleIDPreML::process( + const CalorimeterParticleIDPreML::Input& input, + const CalorimeterParticleIDPreML::Output& output) const { + + const auto [clusters, cluster_assocs] = input; + auto [feature_tensors, target_tensors] = output; + + edm4eic::MutableTensor feature_tensor = feature_tensors->create(); + feature_tensor.addToShape(clusters->size()); + feature_tensor.addToShape(11); // p, E/p, azimuthal, polar, 7 shape parameters + feature_tensor.setElementType(1); // 1 - float + + edm4eic::MutableTensor target_tensor; + if (cluster_assocs) { + target_tensor = target_tensors->create(); + target_tensor.addToShape(clusters->size()); + target_tensor.addToShape(2); // is electron, is hadron + target_tensor.setElementType(7); // 7 - int64 + } + + for (edm4eic::Cluster cluster : *clusters) { + double momentum; + { + // FIXME: use track momentum once matching to tracks becomes available + edm4eic::MCRecoClusterParticleAssociation best_assoc; + for (auto assoc : *cluster_assocs) { + if (assoc.getRec() == cluster) { + if ((not best_assoc.isAvailable()) || (assoc.getWeight() > best_assoc.getWeight())) { + best_assoc = assoc; + } + } + } + if (best_assoc.isAvailable()) { + momentum = edm4hep::utils::magnitude(best_assoc.getSim().getMomentum()); + } else { + warning("Can't find association for cluster. Skipping..."); + continue; + } + } + + feature_tensor.addToFloatData(momentum); + feature_tensor.addToFloatData(cluster.getEnergy() / momentum); + auto pos = cluster.getPosition(); + feature_tensor.addToFloatData(edm4hep::utils::anglePolar(pos)); + feature_tensor.addToFloatData(edm4hep::utils::angleAzimuthal(pos)); + for (int par_ix = 0; par_ix < cluster.shapeParameters_size(); par_ix++) { + feature_tensor.addToFloatData(cluster.getShapeParameters(par_ix)); + } + + if (cluster_assocs) { + edm4eic::MCRecoClusterParticleAssociation best_assoc; + for (auto assoc : *cluster_assocs) { + if (assoc.getRec() == cluster) { + if ((not best_assoc.isAvailable()) || (assoc.getWeight() > best_assoc.getWeight())) { + best_assoc = assoc; + } + } + } + int64_t is_electron = 0, is_pion = 0; + if (best_assoc.isAvailable()) { + is_electron = best_assoc.getSim().getPDG() == 11; + is_pion = best_assoc.getSim().getPDG() != 11; + } + target_tensor.addToInt64Data(is_pion); + target_tensor.addToInt64Data(is_electron); + } + } + + size_t expected_num_entries = feature_tensor.getShape(0) * feature_tensor.getShape(1); + if (feature_tensor.floatData_size() != expected_num_entries) { + error("Inconsistent output tensor shape and element count: {} != {}", feature_tensor.floatData_size(), expected_num_entries); + throw std::runtime_error(fmt::format("Inconsistent output tensor shape and element count: {} != {}", feature_tensor.floatData_size(), expected_num_entries)); + } + } + +} // namespace eicrecon +#endif diff --git a/src/algorithms/onnx/CalorimeterParticleIDPreML.h b/src/algorithms/onnx/CalorimeterParticleIDPreML.h new file mode 100644 index 0000000000..bc4b757eb4 --- /dev/null +++ b/src/algorithms/onnx/CalorimeterParticleIDPreML.h @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024 Dmitry Kalinkin + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "algorithms/interfaces/WithPodConfig.h" + +namespace eicrecon { + +using CalorimeterParticleIDPreMLAlgorithm = + algorithms::Algorithm>, + algorithms::Output>>; + +class CalorimeterParticleIDPreML : public CalorimeterParticleIDPreMLAlgorithm, + public WithPodConfig { + +public: + CalorimeterParticleIDPreML(std::string_view name) + : CalorimeterParticleIDPreMLAlgorithm{name, + {"inputClusters"}, + {"outputFeatureTensor", "outputTargetTensor"}, + ""} { + } + + void init() final; + void process(const Input&, const Output&) const final; +}; + +} // namespace eicrecon diff --git a/src/algorithms/onnx/ONNXInference.cc b/src/algorithms/onnx/ONNXInference.cc new file mode 100644 index 0000000000..2fa80cdb7d --- /dev/null +++ b/src/algorithms/onnx/ONNXInference.cc @@ -0,0 +1,188 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2022 - 2024 Wouter Deconinck, Tooba Ali, Dmitry Kalinkin + +#include + +#if EDM4EIC_VERSION_MAJOR >= 8 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ONNXInference.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(); + } + + static bool check_shape_consistency(const std::vector& shape1, const std::vector& shape2) { + if (shape2.size() != shape1.size()) { + return false; + } + for (size_t ix = 0; ix < shape1.size(); ix++) { + if ((shape1[ix] != -1) && (shape2[ix] != -1) && (shape1[ix] != shape2[ix])) { + return false; + } + } + return true; + } + + template + static Ort::Value iters_to_tensor( + typename std::vector::const_iterator data_begin, + typename std::vector::const_iterator data_end, + std::vector::const_iterator shape_begin, + std::vector::const_iterator shape_end + ) { + Ort::MemoryInfo mem_info = + Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); + auto tensor = Ort::Value::CreateTensor(mem_info, const_cast(&*data_begin), data_end - data_begin, &*shape_begin, shape_end - shape_begin); + return tensor; + } + + void ONNXInference::init() { + // onnxruntime setup + m_env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, name().data()); + Ort::SessionOptions session_options; + try { + m_session = Ort::Session(m_env, m_cfg.modelPath.c_str(), session_options); + Ort::AllocatorWithDefaultOptions allocator; + + // print name/shape of inputs + 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()); + debug("\t{} : {}", m_input_names.at(i), print_shape(m_input_shapes.at(i))); + } + + // print name/shape of outputs + debug("Output Node Name/Shape: {}", m_session.GetOutputCount()); + for (std::size_t i = 0; i < m_session.GetOutputCount(); i++) { + m_output_names.emplace_back(m_session.GetOutputNameAllocated(i, allocator).get()); + + if (m_session.GetOutputTypeInfo(i).GetONNXType() != ONNX_TYPE_TENSOR) { + m_output_shapes.emplace_back(); + debug("\t{} : not a tensor", m_output_names.at(i)); + } else { + m_output_shapes.emplace_back(m_session.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape()); + 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(const Ort::Exception& exception) { + error("ONNX error {}", exception.what()); + throw; + } + } + + void ONNXInference::process( + const ONNXInference::Input& input, + const ONNXInference::Output& output) const { + + const auto [in_tensors] = input; + auto [out_tensors] = output; + + // Require valid inputs + if (in_tensors.size() != m_input_names.size()) { + error("The ONNX model requires {} tensors, whereas {} were provided", m_input_names.size(), in_tensors.size()); + throw std::runtime_error(fmt::format("The ONNX model requires {} tensors, whereas {} were provided", m_input_names.size(), in_tensors.size())); + } + + // Prepare input tensor + std::vector input_tensor_values; + std::vector input_tensors; + + for (int ix = 0; ix < m_input_names.size(); ix++) { + edm4eic::Tensor in_tensor = in_tensors[ix]->at(0); + if (in_tensor.getElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + input_tensors.emplace_back(iters_to_tensor( + in_tensor.floatData_begin(), + in_tensor.floatData_end(), + in_tensor.shape_begin(), + in_tensor.shape_end() + )); + } else if (in_tensor.getElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) { + input_tensors.emplace_back(iters_to_tensor( + in_tensor.int64Data_begin(), + in_tensor.int64Data_end(), + in_tensor.shape_begin(), + in_tensor.shape_end() + )); + } + + auto input_shape = input_tensors[ix].GetTensorTypeAndShapeInfo().GetShape(); + std::vector input_expected_shape = m_input_shapes[ix]; + if (!check_shape_consistency(input_shape, input_expected_shape)) { + error("Input tensor shape incorrect {} != {}", print_shape(input_shape), print_shape(input_expected_shape)); + throw std::runtime_error(fmt::format("Input tensor shape incorrect {} != {}", print_shape(input_shape), print_shape(input_expected_shape))); + } + } + + // Attempt inference + std::vector onnx_values; + try { + onnx_values = 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()); + } catch (const Ort::Exception& exception) { + error("Error running model inference: {}", exception.what()); + throw; + } + + try { + for (size_t ix = 0; ix < onnx_values.size(); ix++) { + Ort::Value &onnx_tensor = onnx_values[ix]; + if (!onnx_tensor.IsTensor()) { + error("The output \"{}\" is not a tensor. ONNXType {} is not yet supported. Skipping...", + m_output_names_char[ix], + static_cast(onnx_tensor.GetTypeInfo().GetONNXType())); + continue; + } + auto onnx_tensor_type = onnx_tensor.GetTensorTypeAndShapeInfo(); + edm4eic::MutableTensor out_tensor = out_tensors[ix]->create(); + out_tensor.setElementType(static_cast(onnx_tensor_type.GetElementType())); + size_t num_values = 1; + for (int64_t dim_size : onnx_tensor_type.GetShape()) { + out_tensor.addToShape(dim_size); + num_values *= dim_size; + } + if (onnx_tensor_type.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + auto *data = onnx_tensor.GetTensorMutableData(); + for (size_t value_ix = 0; value_ix < num_values; value_ix++) { + out_tensor.addToFloatData(data[value_ix]); + } + } else if (onnx_tensor_type.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) { + auto *data = onnx_tensor.GetTensorMutableData(); + for (size_t value_ix = 0; value_ix < num_values; value_ix++) { + out_tensor.addToInt64Data(data[value_ix]); + } + } else { + error("Unsupported ONNXTensorElementDataType {}", static_cast(onnx_tensor_type.GetElementType())); + } + } + } catch (const Ort::Exception& exception) { + error("Error running model inference: {}", exception.what()); + throw; + } + } + +} // namespace eicrecon +#endif diff --git a/src/algorithms/onnx/ONNXInference.h b/src/algorithms/onnx/ONNXInference.h new file mode 100644 index 0000000000..8cc8e91232 --- /dev/null +++ b/src/algorithms/onnx/ONNXInference.h @@ -0,0 +1,50 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2022 - 2024 Sylvester Joosten, Dmitry Romanov, Wouter Deconinck, Dmitry Kalinkin + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "algorithms/interfaces/WithPodConfig.h" +#include "algorithms/onnx/ONNXInferenceConfig.h" + +namespace eicrecon { + +using ONNXInferenceAlgorithm = + algorithms::Algorithm>, + algorithms::Output>>; + +class ONNXInference : public ONNXInferenceAlgorithm, + public WithPodConfig { + +public: + ONNXInference(std::string_view name) + : ONNXInferenceAlgorithm{name, + {"inputTensors"}, + {"outputTensors"}, + ""} { + } + + void init() final; + void process(const Input&, const Output&) const final; + +private: + mutable Ort::Env m_env{nullptr}; + 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/ONNXInferenceConfig.h b/src/algorithms/onnx/ONNXInferenceConfig.h new file mode 100644 index 0000000000..a6e98204a1 --- /dev/null +++ b/src/algorithms/onnx/ONNXInferenceConfig.h @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024 Wouter Deconinck, Dmitry Kalinkin + +#pragma once + +#include + +namespace eicrecon { + + struct ONNXInferenceConfig { + + std::string modelPath; + + }; + +} // eicrecon diff --git a/src/detectors/EEMC/CMakeLists.txt b/src/detectors/EEMC/CMakeLists.txt index 90d48a9b0b..7478a5546c 100644 --- a/src/detectors/EEMC/CMakeLists.txt +++ b/src/detectors/EEMC/CMakeLists.txt @@ -15,6 +15,7 @@ plugin_glob_all(${PLUGIN_NAME}) # Find dependencies plugin_add_dd4hep(${PLUGIN_NAME}) plugin_add_event_model(${PLUGIN_NAME}) +plugin_add_onnxruntime(${PLUGIN_NAME}) # Add include directories (works same as target_include_directories) # plugin_include_directories(${PLUGIN_NAME} SYSTEM PUBLIC ... ) diff --git a/src/detectors/EEMC/EEMC.cc b/src/detectors/EEMC/EEMC.cc index 8908c7ac72..035b4ab3d0 100644 --- a/src/detectors/EEMC/EEMC.cc +++ b/src/detectors/EEMC/EEMC.cc @@ -15,8 +15,15 @@ #include "factories/calorimetry/CalorimeterHitDigi_factory.h" #include "factories/calorimetry/CalorimeterHitReco_factory.h" #include "factories/calorimetry/CalorimeterIslandCluster_factory.h" +#if EDM4EIC_VERSION_MAJOR >= 8 +#include "factories/calorimetry/CalorimeterParticleIDPostML_factory.h" +#include "factories/calorimetry/CalorimeterParticleIDPreML_factory.h" +#endif #include "factories/calorimetry/CalorimeterTruthClustering_factory.h" #include "factories/calorimetry/TrackClusterMergeSplitter_factory.h" +#if EDM4EIC_VERSION_MAJOR >= 8 +#include "factories/meta/ONNXInference_factory.h" +#endif extern "C" { void InitPlugin(JApplication *app) { @@ -112,15 +119,24 @@ extern "C" { app->Add( new JOmniFactoryGeneratorT( +#if EDM4EIC_VERSION_MAJOR >= 8 + "EcalEndcapNClustersWithoutPID", +#else "EcalEndcapNClusters", +#endif {"EcalEndcapNIslandProtoClusters", // edm4eic::ProtoClusterCollection #if EDM4EIC_VERSION_MAJOR >= 7 "EcalEndcapNRawHitAssociations"}, // edm4eic::MCRecoCalorimeterHitAssociationCollection #else "EcalEndcapNHits"}, // edm4hep::SimCalorimeterHitCollection #endif +#if EDM4EIC_VERSION_MAJOR >= 8 + {"EcalEndcapNClustersWithoutPID", // edm4eic::Cluster + "EcalEndcapNClusterAssociationsWithoutPID"}, // edm4eic::MCRecoClusterParticleAssociation +#else {"EcalEndcapNClusters", // edm4eic::Cluster "EcalEndcapNClusterAssociations"}, // edm4eic::MCRecoClusterParticleAssociation +#endif { .energyWeight = "log", .sampFrac = 1.0, @@ -150,6 +166,49 @@ extern "C" { ) ); +#if EDM4EIC_VERSION_MAJOR >= 8 + app->Add(new JOmniFactoryGeneratorT( + "EcalEndcapNParticleIDPreML", + { + "EcalEndcapNClustersWithoutPID", + "EcalEndcapNClusterAssociationsWithoutPID", + }, + { + "EcalEndcapNParticleIDInput_features", + "EcalEndcapNParticleIDTarget", + }, + app + )); + app->Add(new JOmniFactoryGeneratorT( + "EcalEndcapNParticleIDInference", + { + "EcalEndcapNParticleIDInput_features", + }, + { + "EcalEndcapNParticleIDOutput_label", + "EcalEndcapNParticleIDOutput_probability_tensor", + }, + { + .modelPath = "calibrations/onnx/EcalEndcapN_pi_rejection.onnx", + }, + app + )); + app->Add(new JOmniFactoryGeneratorT( + "EcalEndcapNParticleIDPostML", + { + "EcalEndcapNClustersWithoutPID", + "EcalEndcapNClusterAssociationsWithoutPID", + "EcalEndcapNParticleIDOutput_probability_tensor", + }, + { + "EcalEndcapNClusters", + "EcalEndcapNClusterAssociations", + "EcalEndcapNClusterParticleIDs", + }, + app + )); +#endif + app->Add( new JOmniFactoryGeneratorT( "EcalEndcapNSplitMergeClusters", diff --git a/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h b/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h new file mode 100644 index 0000000000..51769c13c0 --- /dev/null +++ b/src/factories/calorimetry/CalorimeterParticleIDPostML_factory.h @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024, Dmitry Kalinkin + +#pragma once + +#include "algorithms/onnx/CalorimeterParticleIDPostML.h" +#include "services/algorithms_init/AlgorithmsInit_service.h" +#include "extensions/jana/JOmniFactory.h" + + +namespace eicrecon { + +class CalorimeterParticleIDPostML_factory : public JOmniFactory { + +public: + using AlgoT = eicrecon::CalorimeterParticleIDPostML; +private: + std::unique_ptr m_algo; + + PodioInput m_cluster_input {this}; + PodioInput m_cluster_assoc_input {this}; + PodioInput m_prediction_tensor_input {this}; + + PodioOutput m_cluster_output {this}; + PodioOutput m_cluster_assoc_output {this}; + PodioOutput m_particle_id_output {this}; + +public: + void Configure() { + m_algo = std::make_unique(GetPrefix()); + m_algo->level(static_cast(logger()->level())); + m_algo->applyConfig(config()); + m_algo->init(); + } + + void ChangeRun(int64_t run_number) { + } + + void Process(int64_t run_number, uint64_t event_number) { + m_algo->process({m_cluster_input(), m_cluster_assoc_input(), m_prediction_tensor_input()}, + {m_cluster_output().get(), m_cluster_assoc_output().get(), m_particle_id_output().get()}); + } +}; + +} // eicrecon diff --git a/src/factories/calorimetry/CalorimeterParticleIDPreML_factory.h b/src/factories/calorimetry/CalorimeterParticleIDPreML_factory.h new file mode 100644 index 0000000000..253abf1f69 --- /dev/null +++ b/src/factories/calorimetry/CalorimeterParticleIDPreML_factory.h @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2024, Dmitry Kalinkin + +#pragma once + +#include "algorithms/onnx/CalorimeterParticleIDPreML.h" +#include "services/algorithms_init/AlgorithmsInit_service.h" +#include "extensions/jana/JOmniFactory.h" + + +namespace eicrecon { + +class CalorimeterParticleIDPreML_factory : public JOmniFactory { + +public: + using AlgoT = eicrecon::CalorimeterParticleIDPreML; +private: + std::unique_ptr m_algo; + + PodioInput m_cluster_input {this}; + PodioInput m_cluster_assoc_input {this}; + + PodioOutput m_feature_tensor_output {this}; + PodioOutput m_target_tensor_output {this}; + +public: + void Configure() { + m_algo = std::make_unique(GetPrefix()); + m_algo->level(static_cast(logger()->level())); + m_algo->applyConfig(config()); + m_algo->init(); + } + + void ChangeRun(int64_t run_number) { + } + + void Process(int64_t run_number, uint64_t event_number) { + m_algo->process({m_cluster_input(), m_cluster_assoc_input()}, + {m_feature_tensor_output().get(), m_target_tensor_output().get()}); + } +}; + +} // eicrecon diff --git a/src/factories/meta/ONNXInference_factory.h b/src/factories/meta/ONNXInference_factory.h new file mode 100644 index 0000000000..6837783518 --- /dev/null +++ b/src/factories/meta/ONNXInference_factory.h @@ -0,0 +1,55 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Copyright (C) 2023 - 2024, Wouter Deconinck, Simon Gardener, Dmitry Kalinkin + +#pragma once + +#include "algorithms/onnx/ONNXInference.h" +#include "services/algorithms_init/AlgorithmsInit_service.h" +#include "extensions/jana/JOmniFactory.h" + + +namespace eicrecon { + +class ONNXInference_factory : public JOmniFactory { + +public: + using AlgoT = eicrecon::ONNXInference; +private: + std::unique_ptr m_algo; + + VariadicPodioInput m_input_tensors {this}; + + VariadicPodioOutput m_output_tensors {this}; + + ParameterRef m_modelPath {this, "modelPath", config().modelPath}; + + Service m_algorithmsInit {this}; + +public: + void Configure() { + m_algo = std::make_unique(GetPrefix()); + m_algo->level(static_cast(logger()->level())); + m_algo->applyConfig(config()); + m_algo->init(); + } + + void ChangeRun(int64_t run_number) { + } + + void Process(int64_t run_number, uint64_t event_number) { + std::vector> in_collections; + for (const auto& in_collection : m_input_tensors()) { + in_collections.push_back(gsl::not_null{in_collection}); + } + + std::vector> out_collections; + for (const auto& out_collection : m_output_tensors()) { + out_collections.push_back(gsl::not_null{out_collection.get()}); + } + + m_algo->process(in_collections, + out_collections); + } +}; + +} // eicrecon From bd00e60afb6a04be0c3198e03105e2a552e2c312 Mon Sep 17 00:00:00 2001 From: Dmitry Kalinkin Date: Sun, 1 Sep 2024 21:40:09 -0400 Subject: [PATCH 2/2] disable threadpools (causes deadlocks in onnxruntime::InferenceSession::~InferenceSession()) --- src/algorithms/onnx/InclusiveKinematicsML.cc | 2 ++ src/algorithms/onnx/ONNXInference.cc | 2 ++ 2 files changed, 4 insertions(+) diff --git a/src/algorithms/onnx/InclusiveKinematicsML.cc b/src/algorithms/onnx/InclusiveKinematicsML.cc index 2483d6c629..091e637eff 100644 --- a/src/algorithms/onnx/InclusiveKinematicsML.cc +++ b/src/algorithms/onnx/InclusiveKinematicsML.cc @@ -34,6 +34,8 @@ namespace eicrecon { // onnxruntime setup m_env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "inclusive-kinematics-ml"); Ort::SessionOptions session_options; + session_options.SetInterOpNumThreads(1); + session_options.SetIntraOpNumThreads(1); try { m_session = Ort::Session(m_env, m_cfg.modelPath.c_str(), session_options); diff --git a/src/algorithms/onnx/ONNXInference.cc b/src/algorithms/onnx/ONNXInference.cc index 2fa80cdb7d..00665cd852 100644 --- a/src/algorithms/onnx/ONNXInference.cc +++ b/src/algorithms/onnx/ONNXInference.cc @@ -54,6 +54,8 @@ namespace eicrecon { // onnxruntime setup m_env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, name().data()); Ort::SessionOptions session_options; + session_options.SetInterOpNumThreads(1); + session_options.SetIntraOpNumThreads(1); try { m_session = Ort::Session(m_env, m_cfg.modelPath.c_str(), session_options); Ort::AllocatorWithDefaultOptions allocator;