TFHandler.cxx

Go to the documentation of this file.

///////////////////////////////////////////////////////////////////////
// TFHander
//
// A module to evaluate a tensorflow model in ART.
//
// \author $Author: Micah Groh
////////////////////////////////////////////////////////////////////////

#include  <iostream>
#include  <string>
#include "TensorFlowHandler/TFHandler.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/public/session_options.h"

namespace tensorflow
{
  TFHandler::TFHandler(std::string ModelPath, int CPUlimit):
    fSession(0),
    fModelPath(ModelPath)
  {
    Initialize(CPUlimit);
  }

  TFHandler::~TFHandler()
  {
    if(fSession) delete fSession;
  }

  void TFHandler::Initialize(int CPUlimit)
  {
    // Initalize a session and create a frozen tensorflow graph
    GraphDef graph_def;
    // Force tf to only use a specified number of cores
    tensorflow::SessionOptions options;
    tensorflow::ConfigProto &config = options.config;
    if (CPUlimit > 0) {
        config.set_inter_op_parallelism_threads(CPUlimit);
        config.set_intra_op_parallelism_threads(CPUlimit);
        config.set_use_per_session_threads(false);
    }
                             
    Status status = NewSession(options, &fSession);
    if (!status.ok()) {
      std::cout<<"Error when making tensorflow session: "<< status.ToString() << std::endl;
      return;
    }

    status = ReadBinaryProto(Env::Default(),fModelPath, &graph_def);
    if (!status.ok()){
      std::cout<<"Error when reading model file: "<< status.ToString() << std::endl;
      return;
    }

    status = fSession->Create(graph_def);
    if (!status.ok()){
      std::cout<<"Error when creating tensorflow graph: "<< status.ToString() << std::endl;
      return;
    }

    std::cout<<"Successfully loaded tensorflow graph."<<std::endl;
  }

  std::vector<Tensor> TFHandler::Predict(std::vector<std::pair<std::string,Tensor>> inputs, 
                                         std::vector<std::string> outputLabels)
  {
    std::vector<Tensor> outputs;

    // Go go go!
    Status status = fSession->Run(inputs,outputLabels,{}, &outputs);

    if (!status.ok()){
      std::cout<<"Error when running session: "<< status.ToString() << std::endl;
    }

    return outputs;
  }

}