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 {

     *         std::vector<Ort::Value> inputTensors;
     *         std::vector<float> inputData_1;   // The input data is filled by users, possibly from the event information.
     *         int64_t batchSize = m_onnxTool->getBatchSize(inputData_1.size(), 0);  // The batch size is determined by the input data size to support dynamic batch size.
     *         m_onnxTool->addInput(inputTensors, inputData_1, 0, batchSize);
     *         std::vector<int64_t> inputData_2;  // Some models may have multiple inputs. Add inputs one by one.
     *         int64_t batchSize_2 = m_onnxTool->getBatchSize(inputData_2.size(), 1);
     *         m_onnxTool->addInput(inputTensors, inputData_2, 1, batchSize_2);
     *     ```
     *    2. create output tensors:
     *      ```c++
     *          std::vector<Ort::Value> outputTensors;
     *          std::vector<float> outputData;   // The output data will be filled by the onnx session.
     *          m_onnxTool->addOutput(outputTensors, outputData, 0, batchSize);
     *      ```
     *   3. perform inference:
     *     ```c++
     *        m_onnxTool->inference(inputTensors, outputTensors);
     *    ```
     *   4. Model outputs will be automatically filled to outputData.
     *
     *
     * @author Xiangyang Ju <xju@cern.ch>
     */
    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