ONNXRuntime.h

 
#ifndef TOOLS_ONNXRUNTIME_H
#define TOOLS_ONNXRUNTIME_H
 
#include <cassert>
#include <map>
#include <memory>
#include <string>
#include <vector>
 
#include "onnxruntime_cxx_api.h"
 
namespace ldmx {
namespace ort {
 
typedef std::vector<std::vector<float>> FloatArrays;
 
class ONNXRuntime {
 public:
  ONNXRuntime(const std::string& model_path,
              const ::Ort::SessionOptions* session_options = nullptr);
  ONNXRuntime(const ONNXRuntime&) = delete;
  ONNXRuntime& operator=(const ONNXRuntime&) = delete;
  ~ONNXRuntime() = default;
 
  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;
 
  const std::vector<std::string>& getOutputNames() const;
 
  const std::vector<int64_t>& getOutputShape(
      const std::string& output_name) const;
 
 private:
  static ::Ort::Env env;
  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_;
};
 
}  // namespace ort
}  // namespace ldmx
 
#endif /* TOOLS_ONNXRUNTIME_H_ */