SPInferenceToolBase.h

Go to the documentation of this file.

#ifndef MUONINFERENCETOOLS_SPINFERENCETOOL_H
#define MUONINFERENCETOOLS_SPINFERENCETOOL_H
 
 
#include "MuonInferenceInterfaces/IGraphInferenceTool.h"
#include "MuonInferenceInterfaces/NodeFeatureList.h"
#include "MuonInferenceInterfaces/GraphData.h"
 
#include "AthOnnxInterfaces/IOnnxRuntimeSessionTool.h"
 
#include "MuonSpacePoint/SpacePointContainer.h"
 
#include "AthenaBaseComps/AthAlgTool.h"
#include "StoreGate/ReadHandleKey.h"
#include <onnxruntime_cxx_api.h> // is this somewhere else?
#include "nlohmann/json.hpp"
 
namespace MuonML{
    class SPInferenceToolBase : public extends<AthAlgTool, IGraphInferenceTool> {
        public:
            using base_class::base_class;
 
            StatusCode buildGraph(const EventContext& ctx,
                                    GraphRawData& graphData) const;
 
            StatusCode runInference(GraphRawData& graphData) const;
        
 
        protected:
            StatusCode setupModel();
    
            Ort::Session& model() const;
            SG::ReadHandleKey<MuonR4::SpacePointContainer> m_readKey{this, "ReadSpacePoints", "MuonSpacePoints"};
        private:
            NodeFeatureList m_graphFeatures{};
            ToolHandle<AthOnnx::IOnnxRuntimeSessionTool> m_onnxSessionTool{this, "ModelSession", "" };
            // Cached at initialize: true when the session runs on CUDA.
            bool m_isCuda{false};
            int  m_cudaDeviceId{0};
 
    }; 
 
}
 
#endif