CMS 3D CMS Logo

List of all members | Public Member Functions | Private Attributes | Static Private Attributes
cms::Ort::ONNXRuntime Class Reference

#include <ONNXRuntime.h>

Public Member Functions

const std::vector< std::string > & getOutputNames () const
 
const std::vector< int64_t > & getOutputShape (const std::string &output_name) const
 
 ONNXRuntime (const ONNXRuntime &)=delete
 
 ONNXRuntime (const std::string &model_path, const ::Ort::SessionOptions *session_options=nullptr)
 
ONNXRuntimeoperator= (const ONNXRuntime &)=delete
 
FloatArrays run (const std::vector< std::string > &input_names, FloatArrays &input_values, const std::vector< std::vector< int64_t >> &input_shapes={}, const std::vector< std::string > &output_names={}, int64_t batch_size=1) const
 
 ~ONNXRuntime ()
 

Private Attributes

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

Static Private Attributes

static const ::Ort::Env env_
 

Detailed Description

Definition at line 25 of file ONNXRuntime.h.

Constructor & Destructor Documentation

◆ ONNXRuntime() [1/2]

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

◆ ONNXRuntime() [2/2]

cms::Ort::ONNXRuntime::ONNXRuntime ( const ONNXRuntime )
delete

◆ ~ONNXRuntime()

cms::Ort::ONNXRuntime::~ONNXRuntime ( )

Definition at line 79 of file ONNXRuntime.cc.

79 {}

Member Function Documentation

◆ getOutputNames()

const std::vector< std::string > & cms::Ort::ONNXRuntime::getOutputNames ( ) const

Definition at line 154 of file ONNXRuntime.cc.

154  {
155  if (session_) {
156  return output_node_strings_;
157  } else {
158  throw cms::Exception("RuntimeError") << "Needs to call createSession() first before getting the output names!";
159  }
160  }

References Exception, output_node_strings_, and session_.

Referenced by TrackQuality::setTrackQuality().

◆ getOutputShape()

const std::vector< int64_t > & cms::Ort::ONNXRuntime::getOutputShape ( const std::string &  output_name) const

Definition at line 162 of file ONNXRuntime.cc.

162  {
163  auto iter = output_node_dims_.find(output_name);
164  if (iter == output_node_dims_.end()) {
165  throw cms::Exception("RuntimeError") << "Output name " << output_name << " is invalid!";
166  } else {
167  return iter->second;
168  }
169  }

References Exception, and output_node_dims_.

◆ operator=()

ONNXRuntime& cms::Ort::ONNXRuntime::operator= ( const ONNXRuntime )
delete

◆ run()

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

Definition at line 81 of file ONNXRuntime.cc.

85  {
86  assert(input_names.size() == input_values.size());
87  assert(input_shapes.empty() || input_names.size() == input_shapes.size());
88  assert(batch_size > 0);
89 
90  // create input tensor objects from data values
91  std::vector<Value> input_tensors;
92  auto memory_info = MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
93  for (const auto& name : input_node_strings_) {
94  auto iter = std::find(input_names.begin(), input_names.end(), name);
95  if (iter == input_names.end()) {
96  throw cms::Exception("RuntimeError") << "Input " << name << " is not provided!";
97  }
98  auto input_pos = iter - input_names.begin();
99  auto value = input_values.begin() + input_pos;
100  std::vector<int64_t> input_dims;
101  if (input_shapes.empty()) {
102  input_dims = input_node_dims_.at(name);
103  input_dims[0] = batch_size;
104  } else {
105  input_dims = input_shapes[input_pos];
106  // rely on the given input_shapes to set the batch size
107  }
108  auto expected_len = std::accumulate(input_dims.begin(), input_dims.end(), 1, std::multiplies<int64_t>());
109  if (expected_len != (int64_t)value->size()) {
110  throw cms::Exception("RuntimeError")
111  << "Input array " << name << " has a wrong size of " << value->size() << ", expected " << expected_len;
112  }
113  auto input_tensor =
114  Value::CreateTensor<float>(memory_info, value->data(), value->size(), input_dims.data(), input_dims.size());
115  assert(input_tensor.IsTensor());
116  input_tensors.emplace_back(std::move(input_tensor));
117  }
118 
119  // set output node names; will get all outputs if `output_names` is not provided
120  std::vector<const char*> run_output_node_names;
121  if (output_names.empty()) {
122  run_output_node_names = output_node_names_;
123  } else {
124  for (const auto& name : output_names) {
125  run_output_node_names.push_back(name.c_str());
126  }
127  }
128 
129  // run
130  auto output_tensors = session_->Run(RunOptions{nullptr},
131  input_node_names_.data(),
132  input_tensors.data(),
133  input_tensors.size(),
134  run_output_node_names.data(),
135  run_output_node_names.size());
136 
137  // convert output to floats
139  for (auto& output_tensor : output_tensors) {
140  assert(output_tensor.IsTensor());
141 
142  // get output shape
143  auto tensor_info = output_tensor.GetTensorTypeAndShapeInfo();
144  auto length = tensor_info.GetElementCount();
145 
146  auto floatarr = output_tensor.GetTensorMutableData<float>();
147  outputs.emplace_back(floatarr, floatarr + length);
148  }
149  assert(outputs.size() == run_output_node_names.size());
150 
151  return outputs;
152  }

References cms::cuda::assert(), Exception, spr::find(), pfDeepBoostedJetPreprocessParams_cfi::input_names, input_node_dims_, input_node_names_, input_node_strings_, eostools::move(), Skims_PA_cff::name, output_node_names_, PatBasicFWLiteJetAnalyzer_Selector_cfg::outputs, and session_.

Member Data Documentation

◆ env_

const Env cms::Ort::ONNXRuntime::env_
staticprivate

Definition at line 54 of file ONNXRuntime.h.

◆ input_node_dims_

std::map<std::string, std::vector<int64_t> > cms::Ort::ONNXRuntime::input_node_dims_
private

Definition at line 59 of file ONNXRuntime.h.

Referenced by run().

◆ input_node_names_

std::vector<const char*> cms::Ort::ONNXRuntime::input_node_names_
private

Definition at line 58 of file ONNXRuntime.h.

Referenced by run().

◆ input_node_strings_

std::vector<std::string> cms::Ort::ONNXRuntime::input_node_strings_
private

Definition at line 57 of file ONNXRuntime.h.

Referenced by run().

◆ output_node_dims_

std::map<std::string, std::vector<int64_t> > cms::Ort::ONNXRuntime::output_node_dims_
private

Definition at line 63 of file ONNXRuntime.h.

Referenced by getOutputShape().

◆ output_node_names_

std::vector<const char*> cms::Ort::ONNXRuntime::output_node_names_
private

Definition at line 62 of file ONNXRuntime.h.

Referenced by run().

◆ output_node_strings_

std::vector<std::string> cms::Ort::ONNXRuntime::output_node_strings_
private

Definition at line 61 of file ONNXRuntime.h.

Referenced by getOutputNames().

◆ session_

std::unique_ptr<::Ort::Session> cms::Ort::ONNXRuntime::session_
private

Definition at line 55 of file ONNXRuntime.h.

Referenced by getOutputNames(), and run().

cms::Ort::ONNXRuntime::session_
std::unique_ptr<::Ort::Session > session_
Definition: ONNXRuntime.h:55
PatBasicFWLiteJetAnalyzer_Selector_cfg.outputs
outputs
Definition: PatBasicFWLiteJetAnalyzer_Selector_cfg.py:48
cms::cuda::assert
assert(be >=bs)
spr::find
void find(edm::Handle< EcalRecHitCollection > &hits, DetId thisDet, std::vector< EcalRecHitCollection::const_iterator > &hit, bool debug=false)
Definition: FindCaloHit.cc:19
cms::Ort::ONNXRuntime::input_node_strings_
std::vector< std::string > input_node_strings_
Definition: ONNXRuntime.h:57
cms::Ort::ONNXRuntime::input_node_names_
std::vector< const char * > input_node_names_
Definition: ONNXRuntime.h:58
cms::Ort::ONNXRuntime::output_node_strings_
std::vector< std::string > output_node_strings_
Definition: ONNXRuntime.h:61
cms::Ort::ONNXRuntime::input_node_dims_
std::map< std::string, std::vector< int64_t > > input_node_dims_
Definition: ONNXRuntime.h:59
pfDeepBoostedJetPreprocessParams_cfi.input_names
input_names
Definition: pfDeepBoostedJetPreprocessParams_cfi.py:4
value
Definition: value.py:1
cms::Ort::FloatArrays
std::vector< std::vector< float > > FloatArrays
Definition: ONNXRuntime.h:23
eostools.move
def move(src, dest)
Definition: eostools.py:511
Exception
Definition: hltDiff.cc:245
Skims_PA_cff.name
name
Definition: Skims_PA_cff.py:17
cms::Ort::ONNXRuntime::output_node_names_
std::vector< const char * > output_node_names_
Definition: ONNXRuntime.h:62
cms::Ort::ONNXRuntime::output_node_dims_
std::map< std::string, std::vector< int64_t > > output_node_dims_
Definition: ONNXRuntime.h:63