ONNX.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#pragma once
 
#include <mva/interface/Expert.h>
#include <mva/interface/Teacher.h>
#include <mva/interface/Options.h>
 
#include <onnxruntime/onnxruntime_cxx_api.h>
 
namespace Belle2 {
  namespace MVA {
    class ONNXOptions : public SpecificOptions {
 
    public:
      virtual void load(const boost::property_tree::ptree&) override {}

      virtual void save(boost::property_tree::ptree&) const override {}

      virtual po::options_description getDescription() override
      {
        return po::options_description("ONNX options");
      }

      virtual std::string getMethod() const override { return "ONNX"; }
    };

    class ONNXTeacher : public Teacher {
 
    public:
      ONNXTeacher(const GeneralOptions& general_options,
                  const ONNXOptions&) : Teacher(general_options) {}

      virtual Weightfile train(Dataset&) const override
      {
        return Weightfile();
      }
    };

    class ONNXTensorView {
    public:
      ONNXTensorView(Dataset& dataset, int nOutputs)
        : m_inputShape{1, dataset.getNumberOfFeatures()}, m_outputData(nOutputs),
          m_outputShape{1, nOutputs}, m_memoryInfo(Ort::MemoryInfo::CreateCpu(
                                                     OrtDeviceAllocator, OrtMemTypeCPU)),
          m_inputTensor(Ort::Value::CreateTensor<float>(
                          m_memoryInfo, dataset.m_input.data(), dataset.m_input.size(),
                          m_inputShape.data(), m_inputShape.size())),
          m_outputTensor(Ort::Value::CreateTensor<float>(
                           m_memoryInfo, m_outputData.data(), m_outputData.size(),
                           m_outputShape.data(), m_outputShape.size())) {}

      Ort::Value* inputTensor() { return &m_inputTensor; }

      Ort::Value* outputTensor() { return &m_outputTensor; }

      std::vector<float> outputData() { return m_outputData; }
    private:
      std::vector<int64_t> m_inputShape;

      std::vector<float> m_outputData;

      std::vector<int64_t> m_outputShape;

      Ort::MemoryInfo m_memoryInfo;

      Ort::Value m_inputTensor;

      Ort::Value m_outputTensor;
    };

    class ONNXExpert : public Expert {
    public:
      virtual void load(Weightfile& weightfile) override;

      virtual std::vector<float> apply(Dataset& testData) const override;

      virtual std::vector<std::vector<float>> applyMulticlass(Dataset& test_data) const override;
 
    private:
      void run(ONNXTensorView& view) const;

      Ort::Env m_env;

      Ort::SessionOptions m_sessionOptions;

      std::unique_ptr<Ort::Session> m_session;

      Ort::RunOptions m_runOptions;

      const char* m_inputNames[1] = {"input"};

      const char* m_outputNames[1] = {"output"};
    };
  } // namespace MVA
} // namespace Belle2