TFCSONNXHandler.h

Go to the documentation of this file.

 
#ifndef TFCSONNXHANDLER_H
#define TFCSONNXHANDLER_H
 
// inherits from
#include "ISF_FastCaloSimEvent/VNetworkBase.h"
 
#include <iostream>
 
// ONNX Runtime include(s).
#include <onnxruntime_cxx_api.h>
 
// For reading and writing to root
#include "TFile.h"
#include "TTree.h"
 
// For storing the lambda function
#include <functional>
 
 
class TFCSONNXHandler : public VNetworkBase {
public:
  // Don't lose the default constructor
  using VNetworkBase::VNetworkBase;
 
  explicit TFCSONNXHandler(const std::string &inputFile);
 
  explicit TFCSONNXHandler(const std::vector<char> &bytes);
 
  TFCSONNXHandler(const TFCSONNXHandler &copy_from);
 
  NetworkOutputs compute(NetworkInputs const &inputs) const override;
 
  // Output to a ttree file
  using VNetworkBase::writeNetToTTree;
 
  void writeNetToTTree(TTree &tree) override;
 
  std::vector<std::string> getOutputLayers() const override;
 
  void deleteAllButNet() override;
 
protected:
  virtual void print(std::ostream &strm) const override;
 
  void setupPersistedVariables() override;
 
  void setupNet() override;
 
private:
  std::vector<char> m_bytes;
  std::vector<char>
  getSerializedSession(const std::string& tree_name = m_defaultTreeName);
  std::vector<char> readBytesFromTTree(TTree &tree);
  void writeBytesToTTree(TTree &tree, const std::vector<char> &bytes);
 
  // unique ptr deletes the object when it goes out of scope
  std::unique_ptr<Ort::Session> m_session; 
 
  void readSerializedSession();
 
  std::vector<const char *> m_inputNodeNames; 
 
#if ORT_API_VERSION > 11
 
  std::vector<Ort::AllocatedStringPtr> m_storeInputNodeNames; 
#endif
 
  std::vector<const char *> m_outputNodeNames; 
 
#if ORT_API_VERSION > 11
 
  std::vector<Ort::AllocatedStringPtr> m_storeOutputNodeNames; 
#endif
 
  std::vector<std::vector<int64_t>> m_inputNodeDims; 
 
  std::vector<std::vector<int64_t>> m_outputNodeDims; 
 
  std::vector<int64_t> m_outputNodeSize; 
 
  template <typename Tin, typename Tout>
  NetworkOutputs computeTemplate(NetworkInputs const &input);
 
  std::function<NetworkOutputs(NetworkInputs)>
      m_computeLambda; 
 
  Ort::MemoryInfo m_memoryInfo = Ort::MemoryInfo::CreateCpu(
      OrtArenaAllocator, OrtMemTypeDefault); 
 
  std::vector<std::string> m_outputLayers; 
 
  // For the streamer
  ClassDefOverride(TFCSONNXHandler, 1);
};
 
#endif // TFCSONNXHANDLER_H