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 {
    namespace ONNX {

      class BaseTensor {
      public:
        virtual ~BaseTensor() {}

        virtual Ort::Value createOrtTensor() = 0;
      };

      template <typename T>
      class Tensor : public BaseTensor {
        std::vector<T> m_values;

        std::vector<int64_t> m_shape;

        Ort::MemoryInfo m_memoryInfo;

        size_t sizeFromShape(const std::vector<int64_t>& shape)
        {
          size_t size = 1;
          for (auto n : shape) size *= n;
          return size;
        }

        void checkShapePositive()
        {
          for (auto n : m_shape) {
            if (n < 0) throw std::invalid_argument("All shape dimensions must be positive");
          }
        }
 
      public:
        Tensor(std::vector<int64_t> shape)
          : m_values(sizeFromShape(shape)), m_shape(std::move(shape)),
            m_memoryInfo(
              Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU))
        {
          checkShapePositive();
        }
 

        Tensor(std::vector<T> values, std::vector<int64_t> shape)
          : m_values(std::move(values)), m_shape(std::move(shape)),
            m_memoryInfo(
              Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU))
        {
          if (sizeFromShape(m_shape) != m_values.size()) {
            throw std::length_error(
              "Size of the given values vector (" + std::to_string(m_values.size()) + ") "
              "does not match the product of the shape dimensions (" + std::to_string(sizeFromShape(m_shape)) + ")");
          }
          checkShapePositive();
        }

        static auto make_shared(std::vector<int64_t> shape)
        {
          return std::make_shared<Tensor>(std::move(shape));
        }

        static auto make_shared(std::vector<T> values,
                                std::vector<int64_t> shape)
        {
          return std::make_shared<Tensor>(std::move(values), std::move(shape));
        }

        auto& at(size_t index) { return m_values.at(index); }

        auto& at(std::vector<size_t> index)
        {
          size_t flat_index = 0;
          size_t stride = 1;
          for (int64_t i = m_shape.size() - 1; i >= 0; --i) {
            if (index[i] >= static_cast<size_t>(m_shape[i])) {
              throw std::out_of_range(
                "index " + std::to_string(index[i]) + " is out of bounds for axis "
                + std::to_string(i) + " with size " + std::to_string(m_shape[i]));
            }
            flat_index += index[i] * stride;
            stride *= m_shape[i];
          }
          return at(flat_index);
        }

        void setValues(const std::vector<T>& values)
        {
          if (m_values.size() != values.size()) {
            throw std::length_error(
              "Size of new values vector (" + std::to_string(values.size()) + ") "
              "differs from internal size (" + std::to_string(m_values.size()) + ")");
          }
          m_values = values;
        }

        Ort::Value createOrtTensor()
        {
          return Ort::Value::CreateTensor(m_memoryInfo, m_values.data(),
                                          m_values.size(), m_shape.data(),
                                          m_shape.size());
        }
      };
 

      class Session {
      public:

        Session(const std::string filename);

        void run(const std::map<std::string, std::shared_ptr<BaseTensor>>& inputMap,
                 const std::map<std::string, std::shared_ptr<BaseTensor>>& outputMap);

        void run(const std::vector<const char*>& inputNames,
                 std::vector<Ort::Value>& inputs,
                 const std::vector<const char*>& outputNames,
                 std::vector<Ort::Value>& outputs);

        const Ort::Session& getOrtSession() { return *m_session; }
 
      private:
        Ort::Env m_env;

        Ort::SessionOptions m_sessionOptions;

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

        Ort::RunOptions m_runOptions;
      };
    } // namespace ONNX

    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"; }

      std::string m_outputName;
    };

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

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

    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 configureInputOutputNames();

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

      ONNXOptions m_specific_options;

      std::string m_inputName;

      std::string m_outputName;
    };
  } // namespace MVA
} // namespace Belle2