algorithms/onnx/InclusiveKinematicsML.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022, 2023 Wouter Deconinck, Tooba Ali
0003
0004 #include <fmt/core.h>
0005 #include <onnxruntime_c_api.h>
0006 #include <onnxruntime_cxx_api.h>
0007 #include <algorithm>
0008 #include <cstddef>
0009 #include <exception>
0010 #include <gsl/pointers>
0011 #include <iterator>
0012 #include <ostream>
0013
0014 #include "InclusiveKinematicsML.h"
0015
0016 namespace eicrecon {
0017
0018   static std::string print_shape(const std::vector<std::int64_t>& v) {
0019     std::stringstream ss("");
0020     for (std::size_t i = 0; i < v.size() - 1; i++) ss << v[i] << "x";
0021     ss << v[v.size() - 1];
0022     return ss.str();
0023   }
0024
0025   template <typename T>
0026   Ort::Value vec_to_tensor(std::vector<T>& data, const std::vector<std::int64_t>& shape) {
0027     Ort::MemoryInfo mem_info =
0028         Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
0029     auto tensor = Ort::Value::CreateTensor<T>(mem_info, data.data(), data.size(), shape.data(), shape.size());
0030     return tensor;
0031   }
0032
0033   void InclusiveKinematicsML::init() {
0034     // onnxruntime setup
0035     m_env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, "inclusive-kinematics-ml");
0036     Ort::SessionOptions session_options;
0037     try {
0038       m_session = Ort::Session(m_env, m_cfg.modelPath.c_str(), session_options);
0039
0040       // print name/shape of inputs
0041       Ort::AllocatorWithDefaultOptions allocator;
0042       debug("Input Node Name/Shape:");
0043       for (std::size_t i = 0; i < m_session.GetInputCount(); i++) {
0044         m_input_names.emplace_back(m_session.GetInputNameAllocated(i, allocator).get());
0045         if (m_session.GetInputTypeInfo(i).GetONNXType() == ONNX_TYPE_TENSOR) {
0046           m_input_shapes.emplace_back(m_session.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape());
0047           debug("\t{} : {}", m_input_names.at(i), print_shape(m_input_shapes.at(i)));
0048         } else {
0049           m_input_shapes.emplace_back();
0050           debug("\t{} : not a tensor", m_input_names.at(i));
0051         }
0052       }
0053
0054       // print name/shape of outputs
0055       debug("Output Node Name/Shape:");
0056       for (std::size_t i = 0; i < m_session.GetOutputCount(); i++) {
0057         m_output_names.emplace_back(m_session.GetOutputNameAllocated(i, allocator).get());
0058         if (m_session.GetOutputTypeInfo(i).GetONNXType() == ONNX_TYPE_TENSOR) {
0059           m_output_shapes.emplace_back(m_session.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape());
0060           debug("\t{} : {}", m_output_names.at(i), print_shape(m_output_shapes.at(i)));
0061         } else {
0062           m_output_shapes.emplace_back();
0063           debug("\t{} : not a tensor", m_output_names.at(i));
0064         }
0065       }
0066
0067       // convert names to char*
0068       m_input_names_char.resize(m_input_names.size(), nullptr);
0069       std::transform(std::begin(m_input_names), std::end(m_input_names), std::begin(m_input_names_char),
0070                      [&](const std::string& str) { return str.c_str(); });
0071       m_output_names_char.resize(m_output_names.size(), nullptr);
0072       std::transform(std::begin(m_output_names), std::end(m_output_names), std::begin(m_output_names_char),
0073                      [&](const std::string& str) { return str.c_str(); });
0074
0075     } catch(std::exception& e) {
0076       error(e.what());
0077     }
0078   }
0079
0080   void InclusiveKinematicsML::process(
0081       const InclusiveKinematicsML::Input& input,
0082       const InclusiveKinematicsML::Output& output) const {
0083
0084     const auto [electron, da] = input;
0085     auto [ml] = output;
0086
0087     // Require valid inputs
0088     if (electron->size() == 0 || da->size() == 0) {
0089       debug("skipping because input collections have no entries");
0090       return;
0091     }
0092
0093     // Assume model has 1 input nodes and 1 output node.
0094     if (m_input_names.size() != 1 || m_output_names.size() != 1) {
0095       debug("skipping because model has incorrect input and output size");
0096       return;
0097     }
0098
0099     // Prepare input tensor
0100     std::vector<float> input_tensor_values;
0101     std::vector<Ort::Value> input_tensors;
0102     for (std::size_t i = 0; i < electron->size(); i++) {
0103       input_tensor_values.push_back(electron->at(i).getX());
0104     }
0105     input_tensors.emplace_back(vec_to_tensor<float>(input_tensor_values, m_input_shapes.front()));
0106
0107     // Double-check the dimensions of the input tensor
0108     if (! input_tensors[0].IsTensor() || input_tensors[0].GetTensorTypeAndShapeInfo().GetShape() != m_input_shapes.front()) {
0109       debug("skipping because input tensor shape incorrect");
0110       return;
0111     }
0112
0113     // Attempt inference
0114     try {
0115       auto output_tensors = m_session.Run(Ort::RunOptions{nullptr}, m_input_names_char.data(), input_tensors.data(),
0116                                           m_input_names_char.size(), m_output_names_char.data(), m_output_names_char.size());
0117
0118       // Double-check the dimensions of the output tensors
0119       if (!output_tensors[0].IsTensor() || output_tensors.size() != m_output_names.size()) {
0120         debug("skipping because output tensor shape incorrect");
0121         return;
0122       }
0123
0124       // Convert output tensor
0125       float* output_tensor_data = output_tensors[0].GetTensorMutableData<float>();
0126       auto x  = output_tensor_data[0];
0127       auto kin = ml->create();
0128       kin.setX(x);
0129
0130     } catch (const Ort::Exception& exception) {
0131       error("error running model inference: {}", exception.what());
0132     }
0133   }
0134
0135 } // namespace eicrecon