FastSimulations.h

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
 
 
#ifndef O2_ZDC_FAST_SIMULATIONS_H
#define O2_ZDC_FAST_SIMULATIONS_H
 
#if __has_include(<onnxruntime/core/session/onnxruntime_cxx_api.h>)
#include <onnxruntime/core/session/onnxruntime_cxx_api.h>
#else
#include <onnxruntime_cxx_api.h>
#endif
#include <optional>
#include <mutex>
 
namespace o2::zdc::fastsim
{
class NeuralFastSimulation
{
 public:
  NeuralFastSimulation(const std::string& modelPath,
                       OrtAllocatorType allocatorType,
                       OrtMemType memoryType,
                       int64_t batchSize);
  virtual ~NeuralFastSimulation() = default;
 
  void initRunSession();
 
  virtual bool setInput(std::vector<std::vector<float>>& input) = 0;
  virtual void run() = 0;
 
  virtual const std::vector<Ort::Value>& getResult() = 0;
 
  [[nodiscard]] size_t getBatchSize() const;
 
 protected:
  void setInputOutputData();
  void setTensors(std::vector<std::vector<float>>& input);
 
  std::string mModelPath;
 
  Ort::Env mEnv;
  Ort::Session* mSession = nullptr; // a pointer so that we can set it up dynamically and independently of constructor
  Ort::AllocatorWithDefaultOptions mAllocator;
  Ort::MemoryInfo mMemoryInfo;
 
  std::vector<std::string> mInputNames;
  std::vector<std::string> mOutputNames;
  std::vector<std::vector<int64_t>> mInputShapes;
  int64_t mBatchSize;
 
  std::vector<Ort::Value> mInputTensors;
};
 
class ConditionalModelSimulation : public NeuralFastSimulation
{
 public:
  ConditionalModelSimulation(const std::string& modelPath, int64_t batchSize);
  ~ConditionalModelSimulation() override = default;
 
  bool setInput(std::vector<std::vector<float>>& input) override;
  void run() override;
  const std::vector<Ort::Value>& getResult() override;
 
 private:
  std::vector<Ort::Value> mModelOutput;
};
 
class BatchHandler
{
 public:
  static BatchHandler& getInstance(size_t batchSize);
  std::optional<std::vector<std::vector<float>>> getBatch(const std::vector<float>& input);
 
  BatchHandler(const BatchHandler&) = delete;
  BatchHandler& operator=(const BatchHandler&) = delete;
 
 private:
  explicit BatchHandler(size_t batchSize);
  ~BatchHandler() = default;
 
  std::mutex mMutex;
  std::vector<std::vector<float>> mBatch;
  size_t mBatchSize;
};
 
std::optional<std::pair<std::vector<float>, std::vector<float>>> loadScales(const std::string& path);
 
} // namespace o2::zdc::fastsim
#endif // O2_ZDC_FAST_SIMULATIONS_H