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::ml
38{
39
41{
42
43 public:
44 // Constructors & destructors
46 OrtModel(std::unordered_map<std::string, std::string> optionsMap);
47 void init(std::unordered_map<std::string, std::string> optionsMap);
48 virtual ~OrtModel();
49
50 // General purpose
51 void initOptions(std::unordered_map<std::string, std::string> optionsMap);
52 void initEnvironment();
53 void initSession();
54 void memoryOnDevice(int32_t = 0);
55 bool isInitialized() { return mInitialized; }
56 void resetSession();
57
58 // Getters
59 std::vector<std::vector<int64_t>> getNumInputNodes() const { return mInputShapes; }
60 std::vector<std::vector<int64_t>> getNumOutputNodes() const { return mOutputShapes; }
61 std::vector<std::string> getInputNames() const { return mInputNames; }
62 std::vector<std::string> getOutputNames() const { return mOutputNames; }
63 Ort::SessionOptions* getSessionOptions();
64 Ort::MemoryInfo* getMemoryInfo();
65 Ort::Env* getEnv();
66 int32_t getIntraOpNumThreads() const { return mIntraOpNumThreads; }
67 int32_t getInterOpNumThreads() const { return mInterOpNumThreads; }
68
69 // Setters
70 void setDeviceId(int32_t id) { mDeviceId = id; }
71 void setIO();
72 void setActiveThreads(int threads) { mIntraOpNumThreads = threads; }
73 void setIntraOpNumThreads(int threads)
74 {
75 if (mDeviceType == "CPU") {
76 mIntraOpNumThreads = threads;
77 }
78 }
79 void setInterOpNumThreads(int threads)
80 {
81 if (mDeviceType == "CPU") {
82 mInterOpNumThreads = threads;
83 }
84 }
85 void setEnv(Ort::Env*);
86
87 // Conversion
88 template <class I, class O>
89 std::vector<O> v2v(std::vector<I>&, bool = true);
90
91 // Inferencing
92 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
93 std::vector<O> inference(std::vector<I>&);
94
95 template <class I, class O>
96 std::vector<O> inference(std::vector<std::vector<I>>&);
97
98 template <class I, class O>
99 void inference(I*, int64_t, O*);
100
101 template <class I, class O>
102 void inference(I**, int64_t, O*);
103
104 void release(bool = false);
105
106 private:
107 // ORT variables -> need to be hidden as pImpl
108 struct OrtVariables;
109 std::unique_ptr<OrtVariables> mPImplOrt;
110
111 // Input & Output specifications of the loaded network
112 std::vector<const char*> mInputNamesChar, mOutputNamesChar;
113 std::vector<std::string> mInputNames, mOutputNames;
114 std::vector<std::vector<int64_t>> mInputShapes, mOutputShapes, mInputShapesCopy, mOutputShapesCopy; // Input shapes
115 std::vector<int64_t> mInputSizePerNode, mOutputSizePerNode; // Output shapes
116 int32_t mInputsTotal = 0, mOutputsTotal = 0; // Total number of inputs and outputs
117
118 // Environment settings
119 bool mInitialized = false;
120 std::string mModelPath, mEnvName = "", mDeviceType = "CPU", mThreadAffinity = ""; // device options should be cpu, rocm, migraphx, cuda
121 int32_t mIntraOpNumThreads = 1, mInterOpNumThreads = 1, mDeviceId = -1, mEnableProfiling = 0, mLoggingLevel = 0, mAllocateDeviceMemory = 0, mEnableOptimizations = 0;
122
123 std::string printShape(const std::vector<int64_t>&);
124 std::string printShape(const std::vector<std::vector<int64_t>>&, std::vector<std::string>&);
125};
126
127} // namespace o2::ml
128
129#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