OnnxRuntimeBase.hpp

Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file OnnxRuntimeBase.hpp

// This file is part of the Acts project.
//
// Copyright (C) 2020 CERN for the benefit of the Acts project
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

#pragma once
#include <vector>

#include <Eigen/Dense>
#include <onnxruntime_cxx_api.h>

namespace Acts {

using NetworkBatchInput =
    Eigen::Array<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;

// General class that sets up the ONNX runtime framework for loading an ML model
// and using it for inference.
class OnnxRuntimeBase {
 public:
  OnnxRuntimeBase() = default;

  OnnxRuntimeBase(Ort::Env& env, const char* modelPath);

  ~OnnxRuntimeBase() = default;

  std::vector<float> runONNXInference(
      std::vector<float>& inputTensorValues) const;

  std::vector<std::vector<float>> runONNXInference(
      NetworkBatchInput& inputTensorValues) const;

  std::vector<std::vector<std::vector<float>>> runONNXInferenceMultiOutput(
      NetworkBatchInput& inputTensorValues) const;

 private:
  std::unique_ptr<Ort::Session> m_session;
  std::vector<Ort::AllocatedStringPtr> m_inputNodeNamesAllocated;
  std::vector<const char*> m_inputNodeNames;
  std::vector<int64_t> m_inputNodeDims;
  std::vector<Ort::AllocatedStringPtr> m_outputNodeNamesAllocated;
  std::vector<const char*> m_outputNodeNames;
  std::vector<std::vector<int64_t>> m_outputNodeDims;
};

}  // namespace Acts