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