TFHandler.h
Go to the documentation of this file.
1 ///////////////////////////////////////////////////////////////////////
2 // TFHander
3 //
4 // A module to evaluate a tensorflow model in ART.
5 //
6 // \author $Author: Micah Groh
7 ////////////////////////////////////////////////////////////////////////
8 
9 #ifndef TFHANDLER_H
10 #define TFHANDLER_H
11 
12 #include <vector>
13 #include "tensorflow/core/public/session.h"
14 #include "tensorflow/core/platform/env.h"
15 
16 namespace tensorflow
17 {
18  /// Wrapper for Tensorflow which handles construction and prediction
19  class TFHandler
20  {
21  public:
22  /// Basic constructor, takes path to model pb
23  TFHandler(std::string model, int CPUlimit=1);
24  ~TFHandler();
25  void Initialize(int CPUlimit);
26 
27  std::vector<Tensor> Predict(std::vector<std::pair<std::string,Tensor>> inputs,
28  std::vector<std::string> outputLabels);
29 
30  private:
31  Session* fSession;
33  };
34 }
35 
36 #endif
TFHandler(std::string model, int CPUlimit=1)
Basic constructor, takes path to model pb.
Definition: TFHandler.cxx:18
std::string fModelPath
Definition: TFHandler.h:32
void Initialize(int CPUlimit)
Definition: TFHandler.cxx:30
std::vector< Tensor > Predict(std::vector< std::pair< std::string, Tensor >> inputs, std::vector< std::string > outputLabels)
Definition: TFHandler.cxx:64
Wrapper for Tensorflow which handles construction and prediction.
Definition: TFHandler.h:19
const XML_Char XML_Content * model
Definition: expat.h:151
Session * fSession
Definition: TFHandler.h:31
enum BeamMode string