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 std::string &model_path, const ::Ort::SessionOptions *session_options=nullptr)
	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::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

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

cms::Ort::ONNXRuntime::ONNXRuntime ( const ONNXRuntime &  )

delete

cms::Ort::ONNXRuntime::~ONNXRuntime

(

)

Definition at line 78 of file ONNXRuntime.cc.

{}

Member Function Documentation

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

(

)

const

Definition at line 153 of file ONNXRuntime.cc.

References Exception, output_node_strings_, and session_.

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

const std::vector< int64_t > & cms::Ort::ONNXRuntime::getOutputShape

(

const std::string &

output_name

)

const

Definition at line 161 of file ONNXRuntime.cc.

References Exception, and output_node_dims_.

                                                                                          {
    auto iter = output_node_dims_.find(output_name);
    if (iter == output_node_dims_.end()) {
      throw cms::Exception("RuntimeError") << "Output name " << output_name << " is invalid!";
    } else {
      return iter->second;
    }
  }

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

delete

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 80 of file ONNXRuntime.cc.

References Exception, spr::find(), input_node_dims_, input_node_names_, eostools::move(), dataset::name, output_node_names_, PatBasicFWLiteJetAnalyzer_Selector_cfg::outputs, and session_.

                                                         {
    assert(input_names.size() == input_values.size());
    assert(input_shapes.empty() || input_names.size() == input_shapes.size());
    assert(batch_size > 0);

    // create input tensor objects from data values
    std::vector<Value> input_tensors;
    auto memory_info = MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    for (const auto& name : input_node_strings_) {
      auto iter = std::find(input_names.begin(), input_names.end(), name);
      if (iter == input_names.end()) {
        throw cms::Exception("RuntimeError") << "Input " << name << " is not provided!";
      }
      auto input_pos = iter - input_names.begin();
      auto value = input_values.begin() + input_pos;
      std::vector<int64_t> input_dims;
      if (input_shapes.empty()) {
        input_dims = input_node_dims_.at(name);
        input_dims[0] = batch_size;
      } else {
        input_dims = input_shapes[input_pos];
        // rely on the given input_shapes to set the batch size
      }
      auto expected_len = std::accumulate(input_dims.begin(), input_dims.end(), 1, std::multiplies<int64_t>());
      if (expected_len != (int64_t)value->size()) {
        throw cms::Exception("RuntimeError")
            << "Input array " << name << " has a wrong size of " << value->size() << ", expected " << expected_len;
      }
      auto input_tensor =
          Value::CreateTensor<float>(memory_info, value->data(), value->size(), input_dims.data(), input_dims.size());
      assert(input_tensor.IsTensor());
      input_tensors.emplace_back(std::move(input_tensor));
    }

    // set output node names; will get all outputs if `output_names` is not provided
    std::vector<const char*> run_output_node_names;
    if (output_names.empty()) {
      run_output_node_names = output_node_names_;
    } else {
      for (const auto& name : output_names) {
        run_output_node_names.push_back(name.c_str());
      }
    }

    // run
    auto output_tensors = session_->Run(RunOptions{nullptr},
                                        input_node_names_.data(),
                                        input_tensors.data(),
                                        input_tensors.size(),
                                        run_output_node_names.data(),
                                        run_output_node_names.size());

    // convert output to floats
    FloatArrays outputs;
    for (auto& output_tensor : output_tensors) {
      assert(output_tensor.IsTensor());

      // get output shape
      auto tensor_info = output_tensor.GetTensorTypeAndShapeInfo();
      auto length = tensor_info.GetElementCount();

      auto floatarr = output_tensor.GetTensorMutableData<float>();
      outputs.emplace_back(floatarr, floatarr + length);
    }
    assert(outputs.size() == run_output_node_names.size());

    return outputs;
  }

Member Data Documentation

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

staticprivate

Definition at line 54 of file ONNXRuntime.h.

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().

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

private

Definition at line 58 of file ONNXRuntime.h.

Referenced by run().

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

private

Definition at line 57 of file ONNXRuntime.h.

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().

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

private

Definition at line 62 of file ONNXRuntime.h.

Referenced by run().

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

private

Definition at line 61 of file ONNXRuntime.h.

Referenced by getOutputNames().

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

private

Definition at line 55 of file ONNXRuntime.h.

Referenced by getOutputNames(), and run().