TensorflowHandler.cxx

Go to the documentation of this file.

////////////////////////////////////////////////////////////////////////
/// \file    TensorflowHandler.h
/// \brief   TensorflowHandler for ShowerLID
/// \author  Andrew Vold - voldx034@umn.edu
////////////////////////////////////////////////////////////////////////

#include  <iostream>
#include  <string>
#include "Utilities/func/EnvExpand.h"
#include "ShowerLID/TensorflowHandler.h"
#include "ShowerLID/ShowerLID.h"
#include "messagefacility/MessageLogger/MessageLogger.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/platform/env.h"
#include "TVector3.h"

namespace tensorflow
{
  TensorflowHandler::TensorflowHandler(fhicl::ParameterSet const& pset)
  { 
  }

  //Much of code written by Jim Fleming
  //https://medium.com/jim-fleming/loading-a-tensorflow-graph-with-the-c-api-4caaff88463f
  float TensorflowHandler::Predict(float sliceE,std::map<int,std::vector<float>> featuremap,bool fRHC)
  {

    //Initialize session
    Session* session;
    Status status = NewSession(SessionOptions(), &session);
    if (!status.ok()) {
      std::cout << status.ToString() << "\n";
      return -5.0;
    }

    std::string Model = "/nova/app/users/vold/ART2/ShowerLID/FHC_Model.pb";
    if(fRHC) Model = "/nova/app/users/vold/ART2/ShowerLID/RHC_Model.pb";

    GraphDef graph_def;
    status = ReadBinaryProto(Env::Default(),Model, &graph_def);
    if (!status.ok()){
      std::cout << status.ToString() << "\n";
      return -5.0;
    }

    //Add graph to session
    status = session->Create(graph_def);
    if (!status.ok()) {
      std::cout << status.ToString() << "\n";
      return -5.0;
    }

    //Set inputs and outputs
    Tensor aux_input(DT_FLOAT, TensorShape({1,1}));
    Tensor main_input(DT_FLOAT, TensorShape({1,7,19}));
    auto Slc = aux_input.tensor<float,2>();
    auto Shw = main_input.tensor<float,3>();
  
    Slc(0,0) = sliceE; 

    int num1 = 0;
    for(std::map<int,std::vector<float>>::iterator i = featuremap.begin(); i != featuremap.end(); i++){
      std::vector<float> vec = i->second;
      int num2 = 0;
      for(std::vector<float>::iterator j = vec.begin(); j != vec.end(); j++){
        Shw(0,num1,num2) = *j;
        num2++;
      }
      num1++;
    }
/*
    for(int i = 0; i < 7; i++){
      for(int j = 0; j < 19; j++){
        Shw(0,i,j) = featuremap[i][j];
      }
    }
*/
    std::vector<std::pair<string,Tensor>> inputs = {
      { "main_input", main_input },
      { "aux_input", aux_input },
    };

    //Initialize outputs
    std::vector<Tensor> outputs;

    //Run session
    status = session->Run(inputs,{"output_node0"},{},&outputs);
    if (!status.ok()){
      std::cout << status.ToString() << "\n";
      return -5.0;
    }

    //Graph output and convert output to scalar
    auto output_c = outputs[0].scalar<float>();

    //Print for debugging purposes
//    std::cout << outputs[0].DebugString() << "\n";
//    std::cout << output_c() << "\n";

    //Free up resources
    session->Close();

    return output_c();

  }//End Predict

}//End namespace slid