File indexing completed on 2024-09-24 22:51:34
0001
0002
0003
0004
0005
0006
0007
0008 #ifndef PHYSICSTOOLS_TENSORFLOW_TENSORFLOW_H
0009 #define PHYSICSTOOLS_TENSORFLOW_TENSORFLOW_H
0010
0011 #include "tensorflow/core/framework/tensor.h"
0012 #include "tensorflow/core/lib/core/threadpool.h"
0013 #include "tensorflow/core/lib/io/path.h"
0014 #include "tensorflow/core/public/session.h"
0015 #include "tensorflow/core/util/tensor_bundle/naming.h"
0016 #include "tensorflow/cc/client/client_session.h"
0017 #include "tensorflow/cc/saved_model/loader.h"
0018 #include "tensorflow/cc/saved_model/constants.h"
0019 #include "tensorflow/cc/saved_model/tag_constants.h"
0020
0021 #include "PhysicsTools/TensorFlow/interface/NoThreadPool.h"
0022 #include "PhysicsTools/TensorFlow/interface/TBBThreadPool.h"
0023
0024 #include "FWCore/Utilities/interface/Exception.h"
0025
0026 namespace tensorflow {
0027
0028 enum class Backend { cpu, cuda, rocm, intel, best };
0029
0030 typedef std::pair<std::string, Tensor> NamedTensor;
0031 typedef std::vector<NamedTensor> NamedTensorList;
0032
0033 struct Options {
0034 int _nThreads;
0035 Backend _backend;
0036 SessionOptions _options;
0037
0038 Options(Backend backend) : _nThreads{1}, _backend{backend} {
0039 setThreading(_nThreads);
0040 setBackend(_backend);
0041 };
0042
0043 Options() : _nThreads{1}, _backend{Backend::cpu} {
0044 setThreading(_nThreads);
0045 setBackend(_backend);
0046 };
0047
0048
0049 void setThreading(int nThreads = 1);
0050
0051
0052
0053 void setBackend(Backend backend = Backend::cpu);
0054
0055 SessionOptions& getSessionOptions() { return _options; };
0056 int getNThreads() const { return _nThreads; };
0057 Backend getBackend() const { return _backend; };
0058 };
0059
0060
0061
0062
0063 MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag = kSavedModelTagServe);
0064
0065
0066
0067
0068 MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag, Options& options);
0069
0070
0071 MetaGraphDef* loadMetaGraph(const std::string& exportDir, const std::string& tag, Options& Options);
0072
0073
0074
0075 GraphDef* loadGraphDef(const std::string& pbFile);
0076
0077
0078 Session* createSession();
0079
0080
0081
0082 Session* createSession(Options& options);
0083
0084
0085
0086
0087
0088 Session* createSession(const MetaGraphDef* metaGraphDef, const std::string& exportDir, Options& options);
0089
0090
0091
0092
0093 Session* createSession(const GraphDef* graphDef);
0094
0095
0096
0097
0098 Session* createSession(const GraphDef* graphDef, Options& options);
0099
0100
0101 bool closeSession(Session*& session);
0102
0103
0104 bool closeSession(const Session*& session);
0105
0106 bool checkEmptyInputs(const NamedTensorList& inputs);
0107
0108
0109
0110
0111
0112 void run(Session* session,
0113 const NamedTensorList& inputs,
0114 const std::vector<std::string>& outputNames,
0115 std::vector<Tensor>* outputs,
0116 const thread::ThreadPoolOptions& threadPoolOptions);
0117
0118
0119 inline void run(const Session* session,
0120 const NamedTensorList& inputs,
0121 const std::vector<std::string>& outputNames,
0122 std::vector<Tensor>* outputs,
0123 const thread::ThreadPoolOptions& threadPoolOptions) {
0124
0125
0126 run(const_cast<Session*>(session), inputs, outputNames, outputs, threadPoolOptions);
0127 }
0128
0129
0130
0131
0132 void run(Session* session,
0133 const NamedTensorList& inputs,
0134 const std::vector<std::string>& outputNames,
0135 std::vector<Tensor>* outputs,
0136 thread::ThreadPoolInterface* threadPool);
0137
0138
0139 inline void run(const Session* session,
0140 const NamedTensorList& inputs,
0141 const std::vector<std::string>& outputNames,
0142 std::vector<Tensor>* outputs,
0143 thread::ThreadPoolInterface* threadPool) {
0144
0145
0146 run(const_cast<Session*>(session), inputs, outputNames, outputs, threadPool);
0147 }
0148
0149
0150
0151
0152 void run(Session* session,
0153 const NamedTensorList& inputs,
0154 const std::vector<std::string>& outputNames,
0155 std::vector<Tensor>* outputs,
0156 const std::string& threadPoolName = "no_threads");
0157
0158
0159 inline void run(const Session* session,
0160 const NamedTensorList& inputs,
0161 const std::vector<std::string>& outputNames,
0162 std::vector<Tensor>* outputs,
0163 const std::string& threadPoolName = "no_threads") {
0164
0165
0166 run(const_cast<Session*>(session), inputs, outputNames, outputs, threadPoolName);
0167 }
0168
0169
0170
0171
0172 void run(Session* session,
0173 const std::vector<std::string>& outputNames,
0174 std::vector<Tensor>* outputs,
0175 const std::string& threadPoolName = "no_threads");
0176
0177
0178 inline void run(const Session* session,
0179 const std::vector<std::string>& outputNames,
0180 std::vector<Tensor>* outputs,
0181 const std::string& threadPoolName = "no_threads") {
0182
0183
0184 run(const_cast<Session*>(session), outputNames, outputs, threadPoolName);
0185 }
0186
0187
0188
0189 struct SessionCache {
0190 std::atomic<GraphDef*> graph;
0191 std::atomic<Session*> session;
0192
0193
0194 SessionCache() {}
0195
0196
0197 template <typename... Args>
0198 SessionCache(const std::string& graphPath, Args&&... sessionArgs) {
0199 createSession(graphPath, std::forward<Args>(sessionArgs)...);
0200 }
0201
0202
0203 ~SessionCache() { closeSession(); }
0204
0205
0206
0207 template <typename... Args>
0208 void createSession(const std::string& graphPath, Args&&... sessionArgs) {
0209 graph.store(loadGraphDef(graphPath));
0210 session.store(tensorflow::createSession(graph.load(), std::forward<Args>(sessionArgs)...));
0211 }
0212
0213
0214 inline const Session* getSession() const { return session.load(); }
0215
0216
0217 void closeSession();
0218 };
0219
0220 }
0221
0222 #endif