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