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
 
ONNXRuntimeoperator= (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 27 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 170 of file ONNXRuntime.cxx.

170 {
171 if (session_) {
172 return output_node_strings_;
173 } else {
174 throw std::runtime_error("ONNXRuntime session is not initialized!");
175 }
176}

◆ 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 178 of file ONNXRuntime.cxx.

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

◆ 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 100 of file ONNXRuntime.cxx.

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

Member Data Documentation

◆ env

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

Definition at line 75 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 80 of file ONNXRuntime.h.

◆ input_node_names_

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

Definition at line 79 of file ONNXRuntime.h.

◆ input_node_strings_

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

Definition at line 78 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 84 of file ONNXRuntime.h.

◆ output_node_names_

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

Definition at line 83 of file ONNXRuntime.h.

◆ output_node_strings_

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

Definition at line 82 of file ONNXRuntime.h.

◆ session_

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

Definition at line 76 of file ONNXRuntime.h.


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