IOnnxRuntimeInferenceTool.h

Go to the documentation of this file.

// Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
#ifndef AthOnnx_IOnnxRuntimeInferenceTool_H
#define AthOnnx_IOnnxRuntimeInferenceTool_H
 
#include "AsgTools/IAsgTool.h"
 
#include <memory>
#include <numeric>
#include <utility>
 
#include <onnxruntime_cxx_api.h>
 
 
namespace AthOnnx {
    class IOnnxRuntimeInferenceTool : virtual public asg::IAsgTool
    {
        ASG_TOOL_INTERFACE(IOnnxRuntimeInferenceTool)
 
        public:
 
        virtual void setBatchSize(int64_t batchSize) = 0;
 
        virtual int64_t getBatchSize(int64_t dataSize, int idx = 0) const = 0;
 
        template <typename T>
        StatusCode addInput(std::vector<Ort::Value>& inputTensors, std::vector<T>& data, unsigned idx = 0, int64_t batchSize = -1) const;
 
        template <typename T>
        StatusCode addOutput(std::vector<Ort::Value>& outputTensors, std::vector<T>& data, unsigned idx = 0, int64_t batchSize = -1) const;
 
 
        virtual StatusCode inference(std::vector<Ort::Value>& inputTensors, std::vector<Ort::Value>& outputTensors) const = 0;
 
        virtual void printModelInfo() const = 0;
 
        protected:
        unsigned m_numInputs;
        unsigned m_numOutputs;
        std::vector<std::vector<int64_t> > m_inputShapes;
        std::vector<std::vector<int64_t> > m_outputShapes;
 
        private:
        template <typename T>
        Ort::Value createTensor(std::vector<T>& data, const std::vector<int64_t>& dataShape, int64_t batchSize) const;
 
    };
 
    #include "IOnnxRuntimeInferenceTool.icc"
} // namespace AthOnnx
 
#endif