OrtNeuralNet.h

Go to the documentation of this file.

#pragma once
 
#include "NeuralNet.h"
 
#include <string>
#include <memory>
#include <functional>
#include <vector>
#include <any>
#include <iostream>
#include <optional>
 
#include <onnxruntime_cxx_api.h>
 
namespace NN
{
    class OrtNeuralNet : public NeuralNet
    {
    public:
        OrtNeuralNet(const std::string& modelPath, bool useCuda = false, int gpuId = 0);
 
        tl::expected<std::vector<Tensor>, NeuralError> infer(const Tensor& input) override;
 
        size_t inputCount() const override;
 
        size_t outputCount() const override;
 
        Shape inputShape(size_t idx = 0) const override;
        
        Shape outputShape(size_t idx = 0) const override;
 
        std::string inputName(size_t idx = 0) const override;
 
        std::string outputName(size_t idx = 0) const override;
        int gpuId() const override;
 
    private:
        Ort::Env m_env;
        Ort::SessionOptions m_sessionOptions;
        Ort::Session m_session{ nullptr };
        Ort::AllocatorWithDefaultOptions m_allocator;
 
        std::vector<int64_t> m_inputShape; // for now only one input is supported!
        std::vector<std::vector<int64_t>> m_outputShapes;
        std::string m_inputName;
        std::vector<std::string> m_outputNames;
        int m_gpuId;
 
        tl::expected<Ort::Value, NeuralError> tensorToOrtValue(const Tensor& tensor, std::optional<std::vector<int64_t>> shapeOverride = std::nullopt) const;
        tl::expected<Tensor, NeuralError> ortValueToTensor(const Ort::Value& value) const;
    };
}