Project
Loading...
Searching...
No Matches
OrtInterface.h
Go to the documentation of this file.
1// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
2// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
3// All rights not expressly granted are reserved.
4//
5// This software is distributed under the terms of the GNU General Public
6// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
7//
8// In applying this license CERN does not waive the privileges and immunities
9// granted to it by virtue of its status as an Intergovernmental Organization
10// or submit itself to any jurisdiction.
11
15
16#ifndef O2_ML_ORTINTERFACE_H
17#define O2_ML_ORTINTERFACE_H
18
19// C++ and system includes
20#include <vector>
21#include <string>
22#include <memory>
23#include <map>
24#include <thread>
25#include <unordered_map>
26
27// O2 includes
28#include "GPUCommonLogger.h"
29
30namespace Ort
31{
32struct SessionOptions;
33struct MemoryInfo;
34struct Env;
35} // namespace Ort
36
37namespace o2
38{
39
40namespace ml
41{
42
44{
45
46 public:
47 // Constructors & destructors
49 OrtModel(std::unordered_map<std::string, std::string> optionsMap);
50 void init(std::unordered_map<std::string, std::string> optionsMap);
51 virtual ~OrtModel();
52
53 // General purpose
54 void initOptions(std::unordered_map<std::string, std::string> optionsMap);
55 void initEnvironment();
56 void initSession();
57 void memoryOnDevice(int32_t = 0);
58 bool isInitialized() { return mInitialized; }
59 void resetSession();
60
61 // Getters
62 std::vector<std::vector<int64_t>> getNumInputNodes() const { return mInputShapes; }
63 std::vector<std::vector<int64_t>> getNumOutputNodes() const { return mOutputShapes; }
64 std::vector<std::string> getInputNames() const { return mInputNames; }
65 std::vector<std::string> getOutputNames() const { return mOutputNames; }
66 Ort::SessionOptions* getSessionOptions();
67 Ort::MemoryInfo* getMemoryInfo();
68 Ort::Env* getEnv();
69 int32_t getIntraOpNumThreads() const { return mIntraOpNumThreads; }
70 int32_t getInterOpNumThreads() const { return mInterOpNumThreads; }
71
72 // Setters
73 void setDeviceId(int32_t id) { mDeviceId = id; }
74 void setIO();
75 void setActiveThreads(int threads) { mIntraOpNumThreads = threads; }
76 void setIntraOpNumThreads(int threads)
77 {
78 if (mDeviceType == "CPU") {
79 mIntraOpNumThreads = threads;
80 }
81 }
82 void setInterOpNumThreads(int threads)
83 {
84 if (mDeviceType == "CPU") {
85 mInterOpNumThreads = threads;
86 }
87 }
88 void setEnv(Ort::Env*);
89
90 // Conversion
91 template <class I, class O>
92 std::vector<O> v2v(std::vector<I>&, bool = true);
93
94 // Inferencing
95 template <class I, class O> // class I is the input data type, e.g. float, class O is the output data type, e.g. OrtDataType::Float16_t from O2/Common/ML/include/ML/GPUORTFloat16.h
96 std::vector<O> inference(std::vector<I>&);
97
98 template <class I, class O>
99 std::vector<O> inference(std::vector<std::vector<I>>&);
100
101 template <class I, class O>
102 void inference(I*, int64_t, O*);
103
104 template <class I, class O>
105 void inference(I**, int64_t, O*);
106
107 void release(bool = false);
108
109 private:
110 // ORT variables -> need to be hidden as pImpl
111 struct OrtVariables;
112 std::unique_ptr<OrtVariables> mPImplOrt;
113
114 // Input & Output specifications of the loaded network
115 std::vector<const char*> mInputNamesChar, mOutputNamesChar;
116 std::vector<std::string> mInputNames, mOutputNames;
117 std::vector<std::vector<int64_t>> mInputShapes, mOutputShapes, mInputShapesCopy, mOutputShapesCopy; // Input shapes
118 std::vector<int64_t> mInputSizePerNode, mOutputSizePerNode; // Output shapes
119 int32_t mInputsTotal = 0, mOutputsTotal = 0; // Total number of inputs and outputs
120
121 // Environment settings
122 bool mInitialized = false;
123 std::string mModelPath, mEnvName = "", mDeviceType = "CPU", mThreadAffinity = ""; // device options should be cpu, rocm, migraphx, cuda
124 int32_t mIntraOpNumThreads = 1, mInterOpNumThreads = 1, mDeviceId = -1, mEnableProfiling = 0, mLoggingLevel = 0, mAllocateDeviceMemory = 0, mEnableOptimizations = 0;
125
126 std::string printShape(const std::vector<int64_t>&);
127 std::string printShape(const std::vector<std::vector<int64_t>>&, std::vector<std::string>&);
128};
129
130} // namespace ml
131
132} // namespace o2
133
134#endif // O2_ML_ORTINTERFACE_H
std::vector< std::vector< int64_t > > getNumInputNodes() const
void initOptions(std::unordered_map< std::string, std::string > optionsMap)
void memoryOnDevice(int32_t=0)
void setActiveThreads(int threads)
Ort::Env * getEnv()
void release(bool=false)
void setEnv(Ort::Env *)
int32_t getInterOpNumThreads() const
std::vector< O > v2v(std::vector< I > &, bool=true)
void setInterOpNumThreads(int threads)
void setIntraOpNumThreads(int threads)
std::vector< std::string > getOutputNames() const
std::vector< std::vector< int64_t > > getNumOutputNodes() const
Ort::MemoryInfo * getMemoryInfo()
void setDeviceId(int32_t id)
std::vector< std::string > getInputNames() const
std::vector< O > inference(std::vector< I > &)
int32_t getIntraOpNumThreads() const
virtual ~OrtModel()
void init(std::unordered_map< std::string, std::string > optionsMap)
Ort::SessionOptions * getSessionOptions()
GLuint id
Definition glcorearb.h:650
a couple of static helper functions to create timestamp values for CCDB queries or override obsolete ...