ONNXWrapper Class Reference

#include <ONNXWrapper.h>

Collaboration diagram for ONNXWrapper:

Public Member Functions
	ONNXWrapper (const std::string &model_path)
std::map< std::string, std::vector< float > >	Run (std::map< std::string, std::vector< float > > inputs, int n_batches=1)
const std::map< std::string, std::vector< int64_t > >	GetModelInputs ()
const std::map< std::string, std::vector< int64_t > >	GetModelOutputs ()
const std::map< std::string, std::string >	GetMETAData ()
std::string	GetMETADataByKey (const char *key)
const std::vector< int64_t > &	getInputShape (int input_nr)
const std::vector< int64_t > &	getOutputShape (int output_nr)
const std::vector< const char * > &	getInputNames ()
const std::vector< const char * > &	getOutputNames ()
int	getNumInputs () const
int	getNumOutputs () const

Private Member Functions
const std::vector< int64_t >	getShape (Ort::TypeInfo model_info)

Private Attributes
std::string	m_modelName
std::string	m_modelPath
size_t	m_nr_inputs
size_t	m_nr_output
std::map< std::string, std::vector< int64_t > >	m_input_dims
std::map< std::string, std::vector< int64_t > >	m_output_dims
std::unique_ptr< Ort::Session >	m_onnxSession
std::unique_ptr< Ort::Env >	m_onnxEnv
Ort::SessionOptions	m_session_options
Ort::AllocatorWithDefaultOptions	m_allocator
std::vector< const char * >	m_output_names
std::vector< const char * >	m_input_names

Detailed Description

Definition at line 20 of file ONNXWrapper.h.

Constructor & Destructor Documentation

ONNXWrapper()

ONNXWrapper::ONNXWrapper

(

const std::string &

model_path

)

Definition at line 9 of file ONNXWrapper.cxx.

                                                    :
  m_modelPath(PathResolverFindCalibFile(model_path)),
  m_onnxEnv(std::make_unique< Ort::Env >(ORT_LOGGING_LEVEL_WARNING, "")) {
   
    // initialize session options if needed
    m_session_options.SetIntraOpNumThreads(1);
    m_session_options.SetGraphOptimizationLevel(
      GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
 
 
    // Initialise the ONNX environment and session using the above options and the model name
    m_onnxSession = std::make_unique<Ort::Session>(*m_onnxEnv,
                                                   m_modelPath.c_str(),
                                                   m_session_options);
 
    // get the input nodes
    m_nr_inputs = m_onnxSession->GetInputCount();
 
    // get the output nodes
    m_nr_output = m_onnxSession->GetOutputCount();
     
 
    // iterate over all input nodes
    for (std::size_t i = 0; i < m_nr_inputs; i++) {
      const char* input_name = m_onnxSession->GetInputNameAllocated(i, m_allocator).release();
 
      m_input_names.push_back(input_name);
 
      m_input_dims[input_name] = getShape(m_onnxSession->GetInputTypeInfo(i));
    }
 
    // iterate over all output nodes
    for(std::size_t i = 0; i < m_nr_output; i++ ) {
      const char* output_name = m_onnxSession->GetOutputNameAllocated(i, m_allocator).release();
 
      m_output_names.push_back(output_name);
 
      m_output_dims[output_name] = getShape(m_onnxSession->GetOutputTypeInfo(i));
 
    }
}

Member Function Documentation

getInputNames()

const std::vector< const char * > & ONNXWrapper::getInputNames

(

)

Definition at line 156 of file ONNXWrapper.cxx.

                                                        {
  //put the model access for input here
  return m_input_names;
}

getInputShape()

const std::vector< int64_t > & ONNXWrapper::getInputShape

(

int

input_nr = 0

)

Definition at line 166 of file ONNXWrapper.cxx.

                                                                  {
  //put the model access for input here
  std::vector<const char*> names = getInputNames();
  return m_input_dims[names.at(input_nr)];
}

GetMETAData()

const std::map< std::string, std::string > ONNXWrapper::GetMETAData

(

)

Definition at line 138 of file ONNXWrapper.cxx.

                                                            {
  std::map<std::string, std::string> METAData_map;
  auto metadata = m_onnxSession->GetModelMetadata();
 
  auto keys = metadata.GetCustomMetadataMapKeysAllocated(m_allocator);
  // auto status = OrtApi::ModelMetadataGetCustomMetadataMapKeys(metadata, m_allocator, keys, nkeys)
  for (size_t i = 0; i < keys.size(); i++) {
    METAData_map[keys[i].get()]=this->GetMETADataByKey(keys[i].get());
  }
 
  return METAData_map;
}

GetMETADataByKey()

std::string ONNXWrapper::GetMETADataByKey

(

const char *

key

)

Definition at line 151 of file ONNXWrapper.cxx.

                                                       {
  auto metadata = m_onnxSession->GetModelMetadata();
  return metadata.LookupCustomMetadataMapAllocated(key, m_allocator).get();
}

GetModelInputs()

const std::map< std::string, std::vector< int64_t > > ONNXWrapper::GetModelInputs

(

)

Definition at line 120 of file ONNXWrapper.cxx.

                                                                        {
  std::map<std::string, std::vector<int64_t>> ModelInputINFO_map;
 
  for(std::size_t i = 0; i < m_nr_inputs; i++ ) {
    ModelInputINFO_map[m_input_names.at(i)] = m_input_dims[m_input_names.at(i)];
  }
  return ModelInputINFO_map;
}

GetModelOutputs()

const std::map< std::string, std::vector< int64_t > > ONNXWrapper::GetModelOutputs

(

)

Definition at line 129 of file ONNXWrapper.cxx.

                                                                         {
  std::map<std::string, std::vector<int64_t>> ModelOutputINFO_map;
 
  for(std::size_t i = 0; i < m_nr_output; i++ ) {
    ModelOutputINFO_map[m_output_names.at(i)] = m_output_dims[m_output_names.at(i)];
  }
  return ModelOutputINFO_map;
}

getNumInputs()

int ONNXWrapper::getNumInputs

(

)

const

Definition at line 178 of file ONNXWrapper.cxx.

{ return m_input_names.size(); }

getNumOutputs()

int ONNXWrapper::getNumOutputs

(

)

const

Definition at line 179 of file ONNXWrapper.cxx.

{ return m_output_names.size(); }

getOutputNames()

const std::vector< const char * > & ONNXWrapper::getOutputNames

(

)

Definition at line 161 of file ONNXWrapper.cxx.

                                                         {
  //put the model access for outputs here
  return m_output_names;
}

getOutputShape()

const std::vector< int64_t > & ONNXWrapper::getOutputShape

(

int

output_nr = 0

)

Definition at line 172 of file ONNXWrapper.cxx.

                                                                    {
  //put the model access for outputs here
  std::vector<const char*> names = getOutputNames();
  return m_output_dims[names.at(output_nr)];
}

getShape()

const std::vector< int64_t > ONNXWrapper::getShape ( Ort::TypeInfo model_info )

private

Definition at line 181 of file ONNXWrapper.cxx.

                                                                     {
      auto tensor_info = model_info.GetTensorTypeAndShapeInfo();
      std::vector<int64_t> dims = tensor_info.GetShape();
      dims[0]=1;
      return dims;
    }

Run()

std::map< std::string, std::vector< float > > ONNXWrapper::Run	(	std::map< std::string, std::vector< float > >	inputs,
		int	n_batches = 1 )

Definition at line 51 of file ONNXWrapper.cxx.

                                                             {
    for ( const auto &p : inputs ) // check for valid dimensions between batches and inputs
    {
      uint64_t n=1;
      for (uint64_t i:m_input_dims[p.first])
        {
          n*=i;
        }
    if ( (p.second.size() % n) != 0){
      
      throw std::invalid_argument("For input '"+p.first+"' length not compatible with model. Expect a multiple of "+std::to_string(n)+", got "+std::to_string(p.second.size()));
    }
    if (  p.second.size()!=(n_batches*n)){
      throw std::invalid_argument("Number of batches not compatible with length of vector");
    }
    }
    // Create a CPU tensor to be used as input
    Ort::MemoryInfo memory_info("Cpu", OrtDeviceAllocator, 0, OrtMemTypeDefault);
 
    // define input tensor
    std::vector<Ort::Value> output_tensor;
    std::vector<Ort::Value> input_tensor;
 
    // add the inputs to vector
    for ( const auto &p : m_input_dims )   
    {        
      std::vector<int64_t> in_dims = p.second;
      in_dims.at(0) = n_batches;
      input_tensor.push_back(Ort::Value::CreateTensor<float>(memory_info,
                                                            inputs[p.first].data(),
                                                            inputs[p.first].size(),
                                                            in_dims.data(),
                                                            in_dims.size()));
    }
 
    // init output tensor and fill with zeros
    std::map<std::string, std::vector<float>> outputs;
    for ( const auto &p : m_output_dims ) {
      std::vector<int64_t> out_dims = p.second;
      out_dims.at(0) = n_batches;
      // init output
      int length = 1;
      for(auto i : out_dims){ length*=i; }
      std::vector<float> output(length,0);
      // std::vector<float> output(m_output_dims[i][1], 0.0);
      outputs[p.first] = std::move(output);
      output_tensor.push_back(Ort::Value::CreateTensor<float>(memory_info,
                                                            outputs[p.first].data(),
                                                            outputs[p.first].size(),
                                                            out_dims.data(),
                                                            out_dims.size()));
    }
 
    Ort::Session& session = *m_onnxSession;
 
    // run the model
    session.Run(Ort::RunOptions{nullptr},
                m_input_names.data(),
                input_tensor.data(),
                2,
                m_output_names.data(),
                output_tensor.data(),
                2); 
 
    return outputs;
    }

Member Data Documentation

m_allocator

Ort::AllocatorWithDefaultOptions ONNXWrapper::m_allocator

private

Definition at line 45 of file ONNXWrapper.h.

m_input_dims

std::map<std::string, std::vector<int64_t> > ONNXWrapper::m_input_dims

private

Definition at line 35 of file ONNXWrapper.h.

m_input_names

std::vector<const char*> ONNXWrapper::m_input_names

private

Definition at line 51 of file ONNXWrapper.h.

m_modelName

std::string ONNXWrapper::m_modelName

private

Definition at line 25 of file ONNXWrapper.h.

m_modelPath

std::string ONNXWrapper::m_modelPath

private

Definition at line 26 of file ONNXWrapper.h.

m_nr_inputs

size_t ONNXWrapper::m_nr_inputs

private

Definition at line 31 of file ONNXWrapper.h.

m_nr_output

size_t ONNXWrapper::m_nr_output

private

Definition at line 32 of file ONNXWrapper.h.

m_onnxEnv

std::unique_ptr< Ort::Env > ONNXWrapper::m_onnxEnv

private

Definition at line 41 of file ONNXWrapper.h.

m_onnxSession

std::unique_ptr<Ort::Session> ONNXWrapper::m_onnxSession

private

Definition at line 40 of file ONNXWrapper.h.

m_output_dims

std::map<std::string, std::vector<int64_t> > ONNXWrapper::m_output_dims

private

Definition at line 36 of file ONNXWrapper.h.

m_output_names

std::vector<const char*> ONNXWrapper::m_output_names

private

Definition at line 50 of file ONNXWrapper.h.

m_session_options

Ort::SessionOptions ONNXWrapper::m_session_options

private

Definition at line 44 of file ONNXWrapper.h.

---

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

• ONNXWrapper.h
• ONNXWrapper.cxx