Onnx/src/OnnxRuntimeBase.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #include "Acts/Plugins/Onnx/OnnxRuntimeBase.hpp"
0010
0011 #include <cassert>
0012 #include <stdexcept>
0013
0014 // Parametrized constructor
0015 Acts::OnnxRuntimeBase::OnnxRuntimeBase(Ort::Env& env, const char* modelPath) {
0016   // Set the ONNX runtime session options
0017   Ort::SessionOptions sessionOptions;
0018   // Set graph optimization level
0019   sessionOptions.SetGraphOptimizationLevel(
0020       GraphOptimizationLevel::ORT_ENABLE_BASIC);
0021   // Create the Ort session
0022   m_session = std::make_unique<Ort::Session>(env, modelPath, sessionOptions);
0023   // Default allocator
0024   Ort::AllocatorWithDefaultOptions allocator;
0025
0026   // Get the names of the input nodes of the model
0027   std::size_t numInputNodes = m_session->GetInputCount();
0028   // Iterate over all input nodes and get the name
0029   for (std::size_t i = 0; i < numInputNodes; i++) {
0030     m_inputNodeNamesAllocated.push_back(
0031         m_session->GetInputNameAllocated(i, allocator));
0032     m_inputNodeNames.push_back(m_inputNodeNamesAllocated.back().get());
0033
0034     // Get the dimensions of the input nodes
0035     // Assumes single input
0036     Ort::TypeInfo inputTypeInfo = m_session->GetInputTypeInfo(i);
0037     auto tensorInfo = inputTypeInfo.GetTensorTypeAndShapeInfo();
0038     m_inputNodeDims = tensorInfo.GetShape();
0039   }
0040
0041   // Get the names of the output nodes
0042   std::size_t numOutputNodes = m_session->GetOutputCount();
0043   // Iterate over all output nodes and get the name
0044   for (std::size_t i = 0; i < numOutputNodes; i++) {
0045     m_outputNodeNamesAllocated.push_back(
0046         m_session->GetOutputNameAllocated(i, allocator));
0047     m_outputNodeNames.push_back(m_outputNodeNamesAllocated.back().get());
0048
0049     // Get the dimensions of the output nodes
0050     Ort::TypeInfo outputTypeInfo = m_session->GetOutputTypeInfo(i);
0051     auto tensorInfo = outputTypeInfo.GetTensorTypeAndShapeInfo();
0052     m_outputNodeDims.push_back(tensorInfo.GetShape());
0053   }
0054 }
0055
0056 // Inference function using ONNX runtime for one single entry
0057 std::vector<float> Acts::OnnxRuntimeBase::runONNXInference(
0058     std::vector<float>& inputTensorValues) const {
0059   Acts::NetworkBatchInput vectorInput(1, inputTensorValues.size());
0060   for (std::size_t i = 0; i < inputTensorValues.size(); i++) {
0061     vectorInput(0, i) = inputTensorValues[i];
0062   }
0063   auto vectorOutput = runONNXInference(vectorInput);
0064   return vectorOutput[0];
0065 }
0066
0067 // Inference function using ONNX runtime
0068 // the function assumes that the model has 1 input node and 1 output node
0069 std::vector<std::vector<float>> Acts::OnnxRuntimeBase::runONNXInference(
0070     Acts::NetworkBatchInput& inputTensorValues) const {
0071   return runONNXInferenceMultiOutput(inputTensorValues).front();
0072 }
0073
0074 // Inference function for single-input, multi-output models
0075 std::vector<std::vector<std::vector<float>>>
0076 Acts::OnnxRuntimeBase::runONNXInferenceMultiOutput(
0077     NetworkBatchInput& inputTensorValues) const {
0078   int batchSize = inputTensorValues.rows();
0079   std::vector<std::int64_t> inputNodeDims = m_inputNodeDims;
0080   std::vector<std::vector<std::int64_t>> outputNodeDims = m_outputNodeDims;
0081
0082   // The first dim node should correspond to the batch size
0083   // If it is -1, it is dynamic and should be set to the input size
0084   if (inputNodeDims[0] == -1) {
0085     inputNodeDims[0] = batchSize;
0086   }
0087
0088   bool outputDimsMatch = true;
0089   for (std::vector<std::int64_t>& nodeDim : outputNodeDims) {
0090     if (nodeDim[0] == -1) {
0091       nodeDim[0] = batchSize;
0092     }
0093     outputDimsMatch &= batchSize == 1 || nodeDim[0] == batchSize;
0094   }
0095
0096   if (batchSize != 1 && (inputNodeDims[0] != batchSize || !outputDimsMatch)) {
0097     throw std::runtime_error(
0098         "runONNXInference: batch size doesn't match the input or output node "
0099         "size");
0100   }
0101
0102   // Create input tensor object from data values
0103   Ort::MemoryInfo memoryInfo =
0104       Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
0105   Ort::Value inputTensor = Ort::Value::CreateTensor<float>(
0106       memoryInfo, inputTensorValues.data(), inputTensorValues.size(),
0107       inputNodeDims.data(), inputNodeDims.size());
0108   // Double-check that inputTensor is a Tensor
0109   if (!inputTensor.IsTensor()) {
0110     throw std::runtime_error(
0111         "runONNXInference: conversion of input to Tensor failed. ");
0112   }
0113   // Score model on input tensors, get back output tensors
0114   Ort::RunOptions run_options;
0115   std::vector<Ort::Value> outputTensors =
0116       m_session->Run(run_options, m_inputNodeNames.data(), &inputTensor,
0117                      m_inputNodeNames.size(), m_outputNodeNames.data(),
0118                      m_outputNodeNames.size());
0119
0120   // Double-check that outputTensors contains Tensors and that the count matches
0121   // that of output nodes
0122   if (!outputTensors[0].IsTensor() ||
0123       (outputTensors.size() != m_outputNodeNames.size())) {
0124     throw std::runtime_error(
0125         "runONNXInference: calculation of output failed. ");
0126   }
0127
0128   std::vector<std::vector<std::vector<float>>> multiOutput;
0129
0130   for (std::size_t i_out = 0; i_out < outputTensors.size(); i_out++) {
0131     // Get pointer to output tensor float values
0132     float* outputTensor = outputTensors.at(i_out).GetTensorMutableData<float>();
0133     // Get the output values
0134     std::vector<std::vector<float>> outputTensorValues(
0135         batchSize, std::vector<float>(outputNodeDims.at(i_out)[1], -1));
0136     for (int i = 0; i < outputNodeDims.at(i_out)[0]; i++) {
0137       for (int j = 0; j < ((outputNodeDims.at(i_out).size() > 1)
0138                                ? outputNodeDims.at(i_out)[1]
0139                                : 1);
0140            j++) {
0141         outputTensorValues[i][j] =
0142             outputTensor[i * outputNodeDims.at(i_out)[1] + j];
0143       }
0144     }
0145     multiOutput.push_back(std::move(outputTensorValues));
0146   }
0147   return multiOutput;
0148 }