TensorFlow/interface/TensorFlow.h

0001 /*
0002  * TensorFlow interface helpers.
0003  * For more info, see https://gitlab.cern.ch/mrieger/CMSSW-DNN.
0004  *
0005  * Author: Marcel Rieger
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     // updates the config of sessionOptions so that it uses nThreads
0049     void setThreading(int nThreads = 1);
0050
0051     // Set the backend option cpu/cuda
0052     // The gpu memory is set to "allow_growth" to avoid TF getting all the CUDA memory at once.
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   // loads a meta graph definition saved at exportDir using the SavedModel interface for a tag and
0061   // predefined options
0062   // transfers ownership
0063   MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag = kSavedModelTagServe);
0064
0065   // loads a meta graph definition saved at exportDir using the SavedModel interface for a tag and
0066   // user provided options
0067   // transfers ownership
0068   MetaGraphDef* loadMetaGraphDef(const std::string& exportDir, const std::string& tag, Options& options);
0069
0070   // deprecated in favor of loadMetaGraphDef
0071   MetaGraphDef* loadMetaGraph(const std::string& exportDir, const std::string& tag, Options& Options);
0072
0073   // loads a graph definition saved as a protobuf file at pbFile
0074   // transfers ownership
0075   GraphDef* loadGraphDef(const std::string& pbFile);
0076
0077   // return a new, empty session using the predefined options
0078   Session* createSession();
0079
0080   // return a new, empty session using user provided options
0081   // transfers ownership
0082   Session* createSession(Options& options);
0083
0084   // return a new session that will contain an already loaded meta graph whose exportDir must be
0085   // given in order to load and initialize the variables, sessionOptions are predefined
0086   // an error is thrown when metaGraphDef is a nullptr or when the graph has no nodes
0087   // transfers ownership
0088   Session* createSession(const MetaGraphDef* metaGraphDef, const std::string& exportDir, Options& options);
0089
0090   // return a new session that will contain an already loaded graph def, sessionOptions are predefined
0091   // an error is thrown when graphDef is a nullptr or when the graph has no nodes
0092   // transfers ownership
0093   Session* createSession(const GraphDef* graphDef);
0094
0095   // return a new session that will contain an already loaded graph def, sessionOptions are user defined
0096   // an error is thrown when graphDef is a nullptr or when the graph has no nodes
0097   // transfers ownership
0098   Session* createSession(const GraphDef* graphDef, Options& options);
0099
0100   // closes a session, calls its destructor, resets the pointer, and returns true on success
0101   bool closeSession(Session*& session);
0102
0103   // version of the function above that accepts a const session
0104   bool closeSession(const Session*& session);
0105
0106   bool checkEmptyInputs(const NamedTensorList& inputs);
0107
0108   // run the session with inputs and outputNames, store output tensors, and control the underlying
0109   // thread pool using threadPoolOptions
0110   // used for thread scheduling with custom thread pool options
0111   // throws a cms exception when not successful
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   // version of the function above that accepts a const session
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     // TF takes a non-const session in the run call which is, however, thread-safe and logically
0125     // const, thus const_cast is consistent
0126     run(const_cast<Session*>(session), inputs, outputNames, outputs, threadPoolOptions);
0127   }
0128
0129   // run the session with inputs and outputNames, store output tensors, and control the underlying
0130   // thread pool
0131   // throws a cms exception when not successful
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   // version of the function above that accepts a const session
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     // TF takes a non-const session in the run call which is, however, thread-safe and logically
0145     // const, thus const_cast is consistent
0146     run(const_cast<Session*>(session), inputs, outputNames, outputs, threadPool);
0147   }
0148
0149   // run the session with inputs and outputNames, store output tensors, and control the underlying
0150   // thread pool using a threadPoolName ("no_threads", "tbb", or "tensorflow")
0151   // throws a cms exception when not successful
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   // version of the function above that accepts a const session
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     // TF takes a non-const session in the run call which is, however, thread-safe and logically
0165     // const, thus const_cast is consistent
0166     run(const_cast<Session*>(session), inputs, outputNames, outputs, threadPoolName);
0167   }
0168
0169   // run the session without inputs but only outputNames, store output tensors, and control the
0170   // underlying thread pool using a threadPoolName ("no_threads", "tbb", or "tensorflow")
0171   // throws a cms exception when not successful
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   // version of the function above that accepts a const session
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     // TF takes a non-const session in the run call which is, however, thread-safe and logically
0183     // const, thus const_cast is consistent
0184     run(const_cast<Session*>(session), outputNames, outputs, threadPoolName);
0185   }
0186
0187   // struct that can be used in edm::stream modules for caching a graph and a session instance,
0188   // both made atomic for cases where access is required from multiple threads
0189   struct SessionCache {
0190     std::atomic<GraphDef*> graph;
0191     std::atomic<Session*> session;
0192
0193     // constructor
0194     SessionCache() {}
0195
0196     // initializing constructor, forwarding all arguments to createSession
0197     template <typename... Args>
0198     SessionCache(const std::string& graphPath, Args&&... sessionArgs) {
0199       createSession(graphPath, std::forward<Args>(sessionArgs)...);
0200     }
0201
0202     // destructor
0203     ~SessionCache() { closeSession(); }
0204
0205     // create the internal graph representation from graphPath and the session object, forwarding
0206     // all additional arguments to the central tensorflow::createSession
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     // return a pointer to the const session
0214     inline const Session* getSession() const { return session.load(); }
0215
0216     // closes and removes the session as well as the graph, and sets the atomic members to nullptr's
0217     void closeSession();
0218   };
0219
0220 }  // namespace tensorflow
0221
0222 #endif  // PHYSICSTOOLS_TENSORFLOW_TENSORFLOW_H