LDMX Software
ldmx::ort::ONNXRuntime Class Reference

A convenience wrapper of the ONNXRuntime C++ API. More...

#include <ONNXRuntime.h>

Public Member Functions

 ONNXRuntime (const std::string &model_path, const ::Ort::SessionOptions *session_options=nullptr)
 Class constructor.
 
 ONNXRuntime (const ONNXRuntime &)=delete
 
ONNXRuntime & operator= (const ONNXRuntime &)=delete
 
FloatArrays run (const std::vector< std::string > &input_names, FloatArrays &input_values, const std::vector< std::string > &output_names={}, int64_t batch_size=1) const
 Run model inference and get outputs.
 
const std::vector< std::string > & getOutputNames () const
 Get the names of all the output nodes.
 
const std::vector< int64_t > & getOutputShape (const std::string &output_name) const
 Get the shape of a output node.
 

Private Attributes

std::unique_ptr<::Ort::Session > session_
 
std::vector< std::string > input_node_strings_
 
std::vector< const char * > input_node_names_
 
std::map< std::string, std::vector< int64_t > > input_node_dims_
 
std::vector< std::string > output_node_strings_
 
std::vector< const char * > output_node_names_
 
std::map< std::string, std::vector< int64_t > > output_node_dims_
 

Static Private Attributes

::Ort::Env env
 

Detailed Description

A convenience wrapper of the ONNXRuntime C++ API.

Definition at line 22 of file ONNXRuntime.h.

Constructor & Destructor Documentation

◆ ONNXRuntime()

ldmx::ort::ONNXRuntime::ONNXRuntime ( const std::string & model_path,
const ::Ort::SessionOptions * session_options = nullptr )

Class constructor.

Parameters
model_pathPath to the ONNX model file.
session_optionsConfiguration options of the ONNXRuntime Session. Leave empty to use the default.

Member Function Documentation

◆ getOutputNames()

const std::vector< std::string > & ldmx::ort::ONNXRuntime::getOutputNames ( ) const

Get the names of all the output nodes.

Returns
A list of names of all the output nodes.

Definition at line 174 of file ONNXRuntime.cxx.

174 {
175 if (session_) {
176 return output_node_strings_;
177 } else {
178 throw std::runtime_error("ONNXRuntime session is not initialized!");
179 }
180}

◆ getOutputShape()

const std::vector< int64_t > & ldmx::ort::ONNXRuntime::getOutputShape ( const std::string & output_name) const

Get the shape of a output node.

The 0th dim depends on the batch size, therefore is set to -1.

Parameters
output_nameName of the output node.
Returns
The shape of the output node as a vector of integers.

Definition at line 182 of file ONNXRuntime.cxx.

183 {
184 auto iter = output_node_dims_.find(output_name);
185 if (iter == output_node_dims_.end()) {
186 throw std::runtime_error("Output name '" + output_name + "' is invalid!");
187 } else {
188 return iter->second;
189 }
190}

◆ run()

FloatArrays ldmx::ort::ONNXRuntime::run ( const std::vector< std::string > & input_names,
FloatArrays & input_values,
const std::vector< std::string > & output_names = {},
int64_t batch_size = 1 ) const

Run model inference and get outputs.

Parameters
input_namesList of the names of the input nodes.
input_valuesList of input arrays for each input node. The order of input_values must match input_names.
output_namesNames of the output nodes to get outputs from. Empty list means all output nodes.
batch_sizeNumber of samples in the batch. Each array in input_values must have a shape layout of (batch_size, ...).
Returns
A std::vector<std::vector<float>>, with the order matched to output_names. When output_names is empty, will return all outputs ordered as in getOutputNames().

Definition at line 104 of file ONNXRuntime.cxx.

107 {
108 assert(input_names.size() == input_values.size());
109 assert(batch_size > 0);
110
111 // create input tensor objects from data values
112 std::vector<Value> input_tensors;
113 auto memory_info =
114 MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
115 for (const auto& name : input_node_strings_) {
116 auto iter = std::find(input_names.begin(), input_names.end(), name);
117 if (iter == input_names.end()) {
118 throw std::runtime_error("Input '" + name + "' is not provided!");
119 }
120 auto value = input_values.begin() + (iter - input_names.begin());
121 auto input_dims = input_node_dims_.at(name);
122 if (input_dims.size() > 0) {
123 input_dims[0] = batch_size;
124 }
125 auto expected_len = std::accumulate(input_dims.begin(), input_dims.end(), 1,
126 std::multiplies<int64_t>());
127 if (expected_len != (int64_t)value->size()) {
128 throw std::runtime_error("Input array '" + name +
129 "' has a wrong size of " +
130 std::to_string(value->size()) + ", expected " +
131 std::to_string(expected_len));
132 }
133 auto input_tensor =
134 Value::CreateTensor<float>(memory_info, value->data(), value->size(),
135 input_dims.data(), input_dims.size());
136 assert(input_tensor.IsTensor());
137 input_tensors.emplace_back(std::move(input_tensor));
138 }
139
140 // set output node names; will get all outputs if `output_names` is not
141 // provided
142 std::vector<const char*> run_output_node_names;
143 if (output_names.empty()) {
144 run_output_node_names = output_node_names_;
145 } else {
146 for (const auto& name : output_names) {
147 run_output_node_names.push_back(name.c_str());
148 }
149 }
150
151 // run
152 auto output_tensors =
153 session_->Run(RunOptions{nullptr}, input_node_names_.data(),
154 input_tensors.data(), input_tensors.size(),
155 run_output_node_names.data(), run_output_node_names.size());
156
157 // convert output to floats
158 FloatArrays outputs;
159 for (auto& output_tensor : output_tensors) {
160 assert(output_tensor.IsTensor());
161
162 // get output shape
163 auto tensor_info = output_tensor.GetTensorTypeAndShapeInfo();
164 auto length = tensor_info.GetElementCount();
165
166 auto floatarr = output_tensor.GetTensorMutableData<float>();
167 outputs.emplace_back(floatarr, floatarr + length);
168 }
169 assert(outputs.size() == run_output_node_names.size());
170
171 return outputs;
172}

Member Data Documentation

◆ env

Env ldmx::ort::ONNXRuntime::env
staticprivate

Definition at line 70 of file ONNXRuntime.h.

◆ input_node_dims_

std::map<std::string, std::vector<int64_t> > ldmx::ort::ONNXRuntime::input_node_dims_
private

Definition at line 75 of file ONNXRuntime.h.

◆ input_node_names_

std::vector<const char*> ldmx::ort::ONNXRuntime::input_node_names_
private

Definition at line 74 of file ONNXRuntime.h.

◆ input_node_strings_

std::vector<std::string> ldmx::ort::ONNXRuntime::input_node_strings_
private

Definition at line 73 of file ONNXRuntime.h.

◆ output_node_dims_

std::map<std::string, std::vector<int64_t> > ldmx::ort::ONNXRuntime::output_node_dims_
private

Definition at line 79 of file ONNXRuntime.h.

◆ output_node_names_

std::vector<const char*> ldmx::ort::ONNXRuntime::output_node_names_
private

Definition at line 78 of file ONNXRuntime.h.

◆ output_node_strings_

std::vector<std::string> ldmx::ort::ONNXRuntime::output_node_strings_
private

Definition at line 77 of file ONNXRuntime.h.

◆ session_

std::unique_ptr<::Ort::Session> ldmx::ort::ONNXRuntime::session_
private

Definition at line 71 of file ONNXRuntime.h.


The documentation for this class was generated from the following files: