CMS 3D CMS Logo

List of all members | Public Types | Public Member Functions | Private Attributes
deep_tau::DeepTauCache Class Reference

#include <DeepTauBase.h>

Public Types

using GraphPtr = std::shared_ptr< tensorflow::GraphDef >
 

Public Member Functions

 DeepTauCache (const std::map< std::string, std::string > &graph_names, bool mem_mapped)
 
const tensorflow::GraphDef & getGraph (const std::string &name="") const
 
tensorflow::Session & getSession (const std::string &name="") const
 
 ~DeepTauCache ()
 

Private Attributes

std::map< std::string, GraphPtrgraphs_
 
std::map< std::string, std::unique_ptr< tensorflow::MemmappedEnv > > memmappedEnv_
 
std::map< std::string, tensorflow::Session * > sessions_
 

Detailed Description

Definition at line 48 of file DeepTauBase.h.

Member Typedef Documentation

◆ GraphPtr

using deep_tau::DeepTauCache::GraphPtr = std::shared_ptr<tensorflow::GraphDef>

Definition at line 50 of file DeepTauBase.h.

Constructor & Destructor Documentation

◆ DeepTauCache()

deep_tau::DeepTauCache::DeepTauCache ( const std::map< std::string, std::string > &  graph_names,
bool  mem_mapped 
)

Definition at line 215 of file DeepTauBase.cc.

References tensorflow::createSession(), Exception, HLT_2022v12_cff::graph_file, graphs_, tensorflow::loadGraphDef(), HLT_2022v12_cff::mem_mapped, memmappedEnv_, AlcaSiPixelAliHarvester0T_cff::options, sessions_, tensorflow::setThreading(), and AlCaHLTBitMon_QueryRunRegistry::string.

215  {
216  for (const auto& graph_entry : graph_names) {
217  tensorflow::SessionOptions options;
219 
220  const std::string& entry_name = graph_entry.first;
221  const std::string& graph_file = graph_entry.second;
222  if (mem_mapped) {
223  memmappedEnv_[entry_name] = std::make_unique<tensorflow::MemmappedEnv>(tensorflow::Env::Default());
224  const tensorflow::Status mmap_status = memmappedEnv_.at(entry_name)->InitializeFromFile(graph_file);
225  if (!mmap_status.ok()) {
226  throw cms::Exception("DeepTauCache: unable to initalize memmapped environment for ")
227  << graph_file << ". \n"
228  << mmap_status.ToString();
229  }
230 
231  graphs_[entry_name] = std::make_unique<tensorflow::GraphDef>();
232  const tensorflow::Status load_graph_status =
233  ReadBinaryProto(memmappedEnv_.at(entry_name).get(),
234  tensorflow::MemmappedFileSystem::kMemmappedPackageDefaultGraphDef,
235  graphs_.at(entry_name).get());
236  if (!load_graph_status.ok())
237  throw cms::Exception("DeepTauCache: unable to load graph from ") << graph_file << ". \n"
238  << load_graph_status.ToString();
239 
240  options.config.mutable_graph_options()->mutable_optimizer_options()->set_opt_level(
241  ::tensorflow::OptimizerOptions::L0);
242  options.env = memmappedEnv_.at(entry_name).get();
243 
244  sessions_[entry_name] = tensorflow::createSession(graphs_.at(entry_name).get(), options);
245 
246  } else {
247  graphs_[entry_name].reset(tensorflow::loadGraphDef(graph_file));
248  sessions_[entry_name] = tensorflow::createSession(graphs_.at(entry_name).get(), options);
249  }
250  }
251  }
Session * createSession(SessionOptions &sessionOptions)
Definition: TensorFlow.cc:85
std::map< std::string, tensorflow::Session * > sessions_
Definition: DeepTauBase.h:62
GraphDef * loadGraphDef(const std::string &pbFile)
Definition: TensorFlow.cc:68
void setThreading(SessionOptions &sessionOptions, int nThreads=1)
Definition: TensorFlow.cc:17
std::map< std::string, GraphPtr > graphs_
Definition: DeepTauBase.h:61
std::map< std::string, std::unique_ptr< tensorflow::MemmappedEnv > > memmappedEnv_
Definition: DeepTauBase.h:63

◆ ~DeepTauCache()

deep_tau::DeepTauCache::~DeepTauCache ( )

Definition at line 253 of file DeepTauBase.cc.

References tensorflow::closeSession(), and sessions_.

253  {
254  for (auto& session_entry : sessions_)
255  tensorflow::closeSession(session_entry.second);
256  }
std::map< std::string, tensorflow::Session * > sessions_
Definition: DeepTauBase.h:62
bool closeSession(Session *&session)
Definition: TensorFlow.cc:198

Member Function Documentation

◆ getGraph()

const tensorflow::GraphDef& deep_tau::DeepTauCache::getGraph ( const std::string &  name = "") const
inline

Definition at line 58 of file DeepTauBase.h.

References graphs_, and Skims_PA_cff::name.

58 { return *graphs_.at(name); }
std::map< std::string, GraphPtr > graphs_
Definition: DeepTauBase.h:61

◆ getSession()

tensorflow::Session& deep_tau::DeepTauCache::getSession ( const std::string &  name = "") const
inline

Definition at line 57 of file DeepTauBase.h.

References Skims_PA_cff::name, and sessions_.

Referenced by DeepTauId::getPartialPredictions(), and DeepTauId::getPredictionsV2().

57 { return *sessions_.at(name); }
std::map< std::string, tensorflow::Session * > sessions_
Definition: DeepTauBase.h:62

Member Data Documentation

◆ graphs_

std::map<std::string, GraphPtr> deep_tau::DeepTauCache::graphs_
private

Definition at line 61 of file DeepTauBase.h.

Referenced by DeepTauCache(), and getGraph().

◆ memmappedEnv_

std::map<std::string, std::unique_ptr<tensorflow::MemmappedEnv> > deep_tau::DeepTauCache::memmappedEnv_
private

Definition at line 63 of file DeepTauBase.h.

Referenced by DeepTauCache().

◆ sessions_

std::map<std::string, tensorflow::Session*> deep_tau::DeepTauCache::sessions_
private

Definition at line 62 of file DeepTauBase.h.

Referenced by DeepTauCache(), getSession(), and ~DeepTauCache().