Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2022-04-26 03:14:00

0001 /*
0002  * ONNXRuntime.cc
0003  *
0004  *  Created on: Jun 28, 2019
0005  *      Author: hqu
0006  */
0007 
0008 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0009 
0010 #include "FWCore/Utilities/interface/Exception.h"
0011 #include "FWCore/Utilities/interface/thread_safety_macros.h"
0012 #include <algorithm>
0013 #include <cassert>
0014 #include <functional>
0015 #include <iostream>
0016 #include <memory>
0017 #include <numeric>
0018 
0019 namespace cms::Ort {
0020 
0021   using namespace ::Ort;
0022 
0023   const Env ONNXRuntime::env_(ORT_LOGGING_LEVEL_ERROR, "");
0024 
0025   ONNXRuntime::ONNXRuntime(const std::string& model_path, const SessionOptions* session_options) {
0026     // create session
0027     if (session_options) {
0028       session_ = std::make_unique<Session>(env_, model_path.c_str(), *session_options);
0029     } else {
0030       session_ = std::make_unique<Session>(env_, model_path.c_str(), defaultSessionOptions());
0031     }
0032     AllocatorWithDefaultOptions allocator;
0033 
0034     // get input names and shapes
0035     size_t num_input_nodes = session_->GetInputCount();
0036     input_node_strings_.resize(num_input_nodes);
0037     input_node_names_.resize(num_input_nodes);
0038     input_node_dims_.clear();
0039 
0040     for (size_t i = 0; i < num_input_nodes; i++) {
0041       // get input node names
0042       std::string input_name(session_->GetInputName(i, allocator));
0043       input_node_strings_[i] = input_name;
0044       input_node_names_[i] = input_node_strings_[i].c_str();
0045 
0046       // get input shapes
0047       auto type_info = session_->GetInputTypeInfo(i);
0048       auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
0049       size_t num_dims = tensor_info.GetDimensionsCount();
0050       input_node_dims_[input_name].resize(num_dims);
0051       tensor_info.GetDimensions(input_node_dims_[input_name].data(), num_dims);
0052     }
0053 
0054     size_t num_output_nodes = session_->GetOutputCount();
0055     output_node_strings_.resize(num_output_nodes);
0056     output_node_names_.resize(num_output_nodes);
0057     output_node_dims_.clear();
0058 
0059     for (size_t i = 0; i < num_output_nodes; i++) {
0060       // get output node names
0061       std::string output_name(session_->GetOutputName(i, allocator));
0062       output_node_strings_[i] = output_name;
0063       output_node_names_[i] = output_node_strings_[i].c_str();
0064 
0065       // get output node types
0066       auto type_info = session_->GetOutputTypeInfo(i);
0067       auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
0068       size_t num_dims = tensor_info.GetDimensionsCount();
0069       output_node_dims_[output_name].resize(num_dims);
0070       tensor_info.GetDimensions(output_node_dims_[output_name].data(), num_dims);
0071 
0072       // the 0th dim depends on the batch size
0073       output_node_dims_[output_name].at(0) = -1;
0074     }
0075   }
0076 
0077   ONNXRuntime::~ONNXRuntime() {}
0078 
0079   SessionOptions ONNXRuntime::defaultSessionOptions(Backend backend) {
0080     SessionOptions sess_opts;
0081     sess_opts.SetIntraOpNumThreads(1);
0082     if (backend == Backend::cuda) {
0083       // https://www.onnxruntime.ai/docs/reference/execution-providers/CUDA-ExecutionProvider.html
0084       OrtCUDAProviderOptions options;
0085       sess_opts.AppendExecutionProvider_CUDA(options);
0086     }
0087     return sess_opts;
0088   }
0089 
0090   FloatArrays ONNXRuntime::run(const std::vector<std::string>& input_names,
0091                                FloatArrays& input_values,
0092                                const std::vector<std::vector<int64_t>>& input_shapes,
0093                                const std::vector<std::string>& output_names,
0094                                int64_t batch_size) const {
0095     assert(input_names.size() == input_values.size());
0096     assert(input_shapes.empty() || input_names.size() == input_shapes.size());
0097     assert(batch_size > 0);
0098 
0099     // create input tensor objects from data values
0100     std::vector<Value> input_tensors;
0101     auto memory_info = MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
0102     for (const auto& name : input_node_strings_) {
0103       auto iter = std::find(input_names.begin(), input_names.end(), name);
0104       if (iter == input_names.end()) {
0105         throw cms::Exception("RuntimeError") << "Input " << name << " is not provided!";
0106       }
0107       auto input_pos = iter - input_names.begin();
0108       auto value = input_values.begin() + input_pos;
0109       std::vector<int64_t> input_dims;
0110       if (input_shapes.empty()) {
0111         input_dims = input_node_dims_.at(name);
0112         input_dims[0] = batch_size;
0113       } else {
0114         input_dims = input_shapes[input_pos];
0115         // rely on the given input_shapes to set the batch size
0116         if (input_dims[0] != batch_size) {
0117           throw cms::Exception("RuntimeError") << "The first element of `input_shapes` (" << input_dims[0]
0118                                                << ") does not match the given `batch_size` (" << batch_size << ")";
0119         }
0120       }
0121       auto expected_len = std::accumulate(input_dims.begin(), input_dims.end(), 1, std::multiplies<int64_t>());
0122       if (expected_len != (int64_t)value->size()) {
0123         throw cms::Exception("RuntimeError")
0124             << "Input array " << name << " has a wrong size of " << value->size() << ", expected " << expected_len;
0125       }
0126       auto input_tensor =
0127           Value::CreateTensor<float>(memory_info, value->data(), value->size(), input_dims.data(), input_dims.size());
0128       assert(input_tensor.IsTensor());
0129       input_tensors.emplace_back(std::move(input_tensor));
0130     }
0131 
0132     // set output node names; will get all outputs if `output_names` is not provided
0133     std::vector<const char*> run_output_node_names;
0134     if (output_names.empty()) {
0135       run_output_node_names = output_node_names_;
0136     } else {
0137       for (const auto& name : output_names) {
0138         run_output_node_names.push_back(name.c_str());
0139       }
0140     }
0141 
0142     // run
0143     auto output_tensors = session_->Run(RunOptions{nullptr},
0144                                         input_node_names_.data(),
0145                                         input_tensors.data(),
0146                                         input_tensors.size(),
0147                                         run_output_node_names.data(),
0148                                         run_output_node_names.size());
0149 
0150     // convert output to floats
0151     FloatArrays outputs;
0152     for (auto& output_tensor : output_tensors) {
0153       assert(output_tensor.IsTensor());
0154 
0155       // get output shape
0156       auto tensor_info = output_tensor.GetTensorTypeAndShapeInfo();
0157       auto length = tensor_info.GetElementCount();
0158 
0159       auto floatarr = output_tensor.GetTensorMutableData<float>();
0160       outputs.emplace_back(floatarr, floatarr + length);
0161     }
0162     assert(outputs.size() == run_output_node_names.size());
0163 
0164     return outputs;
0165   }
0166 
0167   const std::vector<std::string>& ONNXRuntime::getOutputNames() const {
0168     if (session_) {
0169       return output_node_strings_;
0170     } else {
0171       throw cms::Exception("RuntimeError") << "Needs to call createSession() first before getting the output names!";
0172     }
0173   }
0174 
0175   const std::vector<int64_t>& ONNXRuntime::getOutputShape(const std::string& output_name) const {
0176     auto iter = output_node_dims_.find(output_name);
0177     if (iter == output_node_dims_.end()) {
0178       throw cms::Exception("RuntimeError") << "Output name " << output_name << " is invalid!";
0179     } else {
0180       return iter->second;
0181     }
0182   }
0183 
0184 } /* namespace cms::Ort */