algorithms/onnx/ONNXInference.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 - 2024 Wouter Deconinck, Tooba Ali, Dmitry Kalinkin
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 <gsl/pointers>
0010 #include <iterator>
0011 #include <sstream>
0012 #include <stdexcept>
0013
0014 #include "ONNXInference.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++) {
0021     ss << v[i] << " x ";
0022   }
0023   ss << v[v.size() - 1];
0024   return ss.str();
0025 }
0026
0027 static bool check_shape_consistency(const std::vector<std::int64_t>& shape1,
0028                                     const std::vector<std::int64_t>& shape2) {
0029   if (shape2.size() != shape1.size()) {
0030     return false;
0031   }
0032   for (std::size_t ix = 0; ix < shape1.size(); ix++) {
0033     if ((shape1[ix] != -1) && (shape2[ix] != -1) && (shape1[ix] != shape2[ix])) {
0034       return false;
0035     }
0036   }
0037   return true;
0038 }
0039
0040 template <typename T>
0041 static Ort::Value iters_to_tensor(typename std::vector<T>::const_iterator data_begin,
0042                                   typename std::vector<T>::const_iterator data_end,
0043                                   std::vector<int64_t>::const_iterator shape_begin,
0044                                   std::vector<int64_t>::const_iterator shape_end) {
0045   Ort::MemoryInfo mem_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator,
0046                                                         OrtMemType::OrtMemTypeDefault);
0047   auto tensor =
0048       Ort::Value::CreateTensor<T>(mem_info, const_cast<T*>(&*data_begin), data_end - data_begin,
0049                                   &*shape_begin, shape_end - shape_begin);
0050   return tensor;
0051 }
0052
0053 void ONNXInference::init() {
0054   // onnxruntime setup
0055   m_env = Ort::Env(ORT_LOGGING_LEVEL_WARNING, name().data());
0056   Ort::SessionOptions session_options;
0057   session_options.SetInterOpNumThreads(1);
0058   session_options.SetIntraOpNumThreads(1);
0059   try {
0060     m_session = Ort::Session(m_env, m_cfg.modelPath.c_str(), session_options);
0061     Ort::AllocatorWithDefaultOptions allocator;
0062
0063     // print name/shape of inputs
0064     debug("Input Node Name/Shape:");
0065     for (std::size_t i = 0; i < m_session.GetInputCount(); i++) {
0066       m_input_names.emplace_back(m_session.GetInputNameAllocated(i, allocator).get());
0067       m_input_shapes.emplace_back(
0068           m_session.GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape());
0069       debug("\t{} : {}", m_input_names.at(i), print_shape(m_input_shapes.at(i)));
0070     }
0071
0072     // print name/shape of outputs
0073     debug("Output Node Name/Shape: {}", m_session.GetOutputCount());
0074     for (std::size_t i = 0; i < m_session.GetOutputCount(); i++) {
0075       m_output_names.emplace_back(m_session.GetOutputNameAllocated(i, allocator).get());
0076
0077       if (m_session.GetOutputTypeInfo(i).GetONNXType() != ONNX_TYPE_TENSOR) {
0078         m_output_shapes.emplace_back();
0079         debug("\t{} : not a tensor", m_output_names.at(i));
0080       } else {
0081         m_output_shapes.emplace_back(
0082             m_session.GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape());
0083         debug("\t{} : {}", m_output_names.at(i), print_shape(m_output_shapes.at(i)));
0084       }
0085     }
0086
0087     // convert names to char*
0088     m_input_names_char.resize(m_input_names.size(), nullptr);
0089     std::ranges::transform(m_input_names, std::begin(m_input_names_char),
0090                            [&](const std::string& str) { return str.c_str(); });
0091     m_output_names_char.resize(m_output_names.size(), nullptr);
0092     std::ranges::transform(m_output_names, std::begin(m_output_names_char),
0093                            [&](const std::string& str) { return str.c_str(); });
0094
0095   } catch (const Ort::Exception& exception) {
0096     error("ONNX error {}", exception.what());
0097     throw;
0098   }
0099 }
0100
0101 void ONNXInference::process(const ONNXInference::Input& input,
0102                             const ONNXInference::Output& output) const {
0103
0104   const auto [in_tensors] = input;
0105   auto [out_tensors]      = output;
0106
0107   // Require valid inputs
0108   if (in_tensors.size() != m_input_names.size()) {
0109     error("The ONNX model requires {} tensors, whereas {} were provided", m_input_names.size(),
0110           in_tensors.size());
0111     throw std::runtime_error(
0112         fmt::format("The ONNX model requires {} tensors, whereas {} were provided",
0113                     m_input_names.size(), in_tensors.size()));
0114   }
0115
0116   // Prepare input tensor
0117   std::vector<float> input_tensor_values;
0118   std::vector<Ort::Value> input_tensors;
0119
0120   for (std::size_t ix = 0; ix < m_input_names.size(); ix++) {
0121     edm4eic::Tensor in_tensor = in_tensors[ix]->at(0);
0122     if (in_tensor.getElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) {
0123       input_tensors.emplace_back(
0124           iters_to_tensor<float>(in_tensor.floatData_begin(), in_tensor.floatData_end(),
0125                                  in_tensor.shape_begin(), in_tensor.shape_end()));
0126     } else if (in_tensor.getElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) {
0127       input_tensors.emplace_back(
0128           iters_to_tensor<int64_t>(in_tensor.int64Data_begin(), in_tensor.int64Data_end(),
0129                                    in_tensor.shape_begin(), in_tensor.shape_end()));
0130     }
0131
0132     auto input_shape = input_tensors[ix].GetTensorTypeAndShapeInfo().GetShape();
0133     std::vector<std::int64_t> input_expected_shape = m_input_shapes[ix];
0134     if (!check_shape_consistency(input_shape, input_expected_shape)) {
0135       error("Input tensor shape incorrect {} != {}", print_shape(input_shape),
0136             print_shape(input_expected_shape));
0137       throw std::runtime_error(fmt::format("Input tensor shape incorrect {} != {}",
0138                                            print_shape(input_shape),
0139                                            print_shape(input_expected_shape)));
0140     }
0141   }
0142
0143   // Attempt inference
0144   std::vector<Ort::Value> onnx_values;
0145   try {
0146     onnx_values = m_session.Run(Ort::RunOptions{nullptr}, m_input_names_char.data(),
0147                                 input_tensors.data(), m_input_names_char.size(),
0148                                 m_output_names_char.data(), m_output_names_char.size());
0149   } catch (const Ort::Exception& exception) {
0150     error("Error running model inference: {}", exception.what());
0151     throw;
0152   }
0153
0154   try {
0155     for (std::size_t ix = 0; ix < onnx_values.size(); ix++) {
0156       Ort::Value& onnx_tensor = onnx_values[ix];
0157       if (!onnx_tensor.IsTensor()) {
0158         error("The output \"{}\" is not a tensor. ONNXType {} is not yet supported. Skipping...",
0159               m_output_names_char[ix], static_cast<int>(onnx_tensor.GetTypeInfo().GetONNXType()));
0160         continue;
0161       }
0162       auto onnx_tensor_type             = onnx_tensor.GetTensorTypeAndShapeInfo();
0163       edm4eic::MutableTensor out_tensor = out_tensors[ix]->create();
0164       out_tensor.setElementType(static_cast<int32_t>(onnx_tensor_type.GetElementType()));
0165       std::size_t num_values = 1;
0166       for (int64_t dim_size : onnx_tensor_type.GetShape()) {
0167         out_tensor.addToShape(dim_size);
0168         num_values *= dim_size;
0169       }
0170       if (onnx_tensor_type.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) {
0171         auto* data = onnx_tensor.GetTensorMutableData<float>();
0172         for (std::size_t value_ix = 0; value_ix < num_values; value_ix++) {
0173           out_tensor.addToFloatData(data[value_ix]);
0174         }
0175       } else if (onnx_tensor_type.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) {
0176         auto* data = onnx_tensor.GetTensorMutableData<int64_t>();
0177         for (std::size_t value_ix = 0; value_ix < num_values; value_ix++) {
0178           out_tensor.addToInt64Data(data[value_ix]);
0179         }
0180       } else {
0181         error("Unsupported ONNXTensorElementDataType {}",
0182               static_cast<int>(onnx_tensor_type.GetElementType()));
0183       }
0184     }
0185   } catch (const Ort::Exception& exception) {
0186     error("Error running model inference: {}", exception.what());
0187     throw;
0188   }
0189 }
0190
0191 } // namespace eicrecon