Merge pull request #40333 from riga/use_global_session_L1Trigger

Move TF session in L1Trigger/L1NNTauProducer to global cache
cms-sw · Dec 19, 2022 · 9318eb0 · 9318eb0
2 parents cdf3761 + ce1a224
commit 9318eb0
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 29 deletions.
diff --git a/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h b/L1Trigger/Phase2L1ParticleFlow/interface/TauNNId.h
@@ -5,22 +5,20 @@
 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
 #include "DataFormats/L1TParticleFlow/interface/PFCandidate.h"
 
-struct TauNNTFCache {
-  TauNNTFCache() : graphDef(nullptr) {}
-  std::atomic<tensorflow::GraphDef *> graphDef;
-};
-
 class TauNNId {
 public:
-  TauNNId(const std::string &iInput, const TauNNTFCache *cache, const std::string &iWeightFile, int iNParticles);
-  ~TauNNId();
+  TauNNId(const std::string &iInput,
+          const tensorflow::Session *session,
+          const std::string &iWeightFile,
+          int iNParticles);
+  ~TauNNId(){};
 
   void setNNVectorVar();
   float EvaluateNN();
   float compute(const l1t::PFCandidate &iSeed, l1t::PFCandidateCollection &iParts);
 
 private:
-  tensorflow::Session *session_;
+  const tensorflow::Session *session_;
   std::vector<float> NNvectorVar_;
   std::string fInput_;
   int fNParticles_;

diff --git a/L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc b/L1Trigger/Phase2L1ParticleFlow/plugins/L1NNTauProducer.cc
@@ -14,14 +14,14 @@
 
 using namespace l1t;
 
-class L1NNTauProducer : public edm::stream::EDProducer<edm::GlobalCache<TauNNTFCache>> {
+class L1NNTauProducer : public edm::stream::EDProducer<edm::GlobalCache<tensorflow::SessionCache>> {
 public:
-  explicit L1NNTauProducer(const edm::ParameterSet&, const TauNNTFCache*);
+  explicit L1NNTauProducer(const edm::ParameterSet&, const tensorflow::SessionCache*);
   ~L1NNTauProducer() override;
 
   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
-  static std::unique_ptr<TauNNTFCache> initializeGlobalCache(const edm::ParameterSet&);
-  static void globalEndJob(const TauNNTFCache*);
+  static std::unique_ptr<tensorflow::SessionCache> initializeGlobalCache(const edm::ParameterSet&);
+  static void globalEndJob(const tensorflow::SessionCache*){};
 
 private:
   std::unique_ptr<TauNNId> fTauNNId_;
@@ -40,31 +40,27 @@ class L1NNTauProducer : public edm::stream::EDProducer<edm::GlobalCache<TauNNTFC
 
 static constexpr float track_trigger_eta_max = 2.5;
 
-L1NNTauProducer::L1NNTauProducer(const edm::ParameterSet& cfg, const TauNNTFCache* cache)
+L1NNTauProducer::L1NNTauProducer(const edm::ParameterSet& cfg, const tensorflow::SessionCache* cache)
     : fSeedPt_(cfg.getParameter<double>("seedpt")),
       fConeSize_(cfg.getParameter<double>("conesize")),
       fTauSize_(cfg.getParameter<double>("tausize")),
       fMaxTaus_(cfg.getParameter<int>("maxtaus")),
       fNParticles_(cfg.getParameter<int>("nparticles")),
       fL1PFToken_(consumes<vector<l1t::PFCandidate>>(cfg.getParameter<edm::InputTag>("L1PFObjects"))) {
   std::string lNNFile = cfg.getParameter<std::string>("NNFileName");  //,"L1Trigger/Phase2L1Taus/data/tau_3layer.pb");
-  fTauNNId_ = std::make_unique<TauNNId>(
-      lNNFile.find("v0") == std::string::npos ? "input_1:0" : "dense_1_input:0", cache, lNNFile, fNParticles_);
+  fTauNNId_ = std::make_unique<TauNNId>(lNNFile.find("v0") == std::string::npos ? "input_1:0" : "dense_1_input:0",
+                                        cache->getSession(),
+                                        lNNFile,
+                                        fNParticles_);
   produces<l1t::PFTauCollection>("L1PFTausNN");
 }
-std::unique_ptr<TauNNTFCache> L1NNTauProducer::initializeGlobalCache(const edm::ParameterSet& cfg) {
+
+std::unique_ptr<tensorflow::SessionCache> L1NNTauProducer::initializeGlobalCache(const edm::ParameterSet& cfg) {
   tensorflow::setLogging("3");
-  std::string lNNFile = cfg.getParameter<std::string>("NNFileName");
-  edm::FileInPath fp(lNNFile);
-  TauNNTFCache* cache = new TauNNTFCache();
-  cache->graphDef = tensorflow::loadGraphDef(fp.fullPath());
-  return std::unique_ptr<TauNNTFCache>(cache);
-}
-void L1NNTauProducer::globalEndJob(const TauNNTFCache* cache) {
-  if (cache->graphDef != nullptr) {
-    delete cache->graphDef;
-  }
+  std::string graphPath = edm::FileInPath(cfg.getParameter<std::string>("NNFileName")).fullPath();
+  return std::make_unique<tensorflow::SessionCache>(graphPath);
 }
+
 void L1NNTauProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
   edm::Handle<l1t::PFCandidateCollection> l1PFCandidates;
   iEvent.getByToken(fL1PFToken_, l1PFCandidates);

diff --git a/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc b/L1Trigger/Phase2L1ParticleFlow/src/TauNNId.cc
@@ -5,10 +5,13 @@
 
 static constexpr unsigned int n_particles_max = 10;
 
-TauNNId::TauNNId(const std::string &iInput, const TauNNTFCache *cache, const std::string &iWeightFile, int iNParticles) {
+TauNNId::TauNNId(const std::string &iInput,
+                 const tensorflow::Session *session,
+                 const std::string &iWeightFile,
+                 int iNParticles)
+    : session_(session) {
   NNvectorVar_.clear();
   edm::FileInPath fp(iWeightFile);
-  session_ = tensorflow::createSession(cache->graphDef);
   fNParticles_ = iNParticles;
 
   fPt_ = std::make_unique<float[]>(fNParticles_);
@@ -18,7 +21,6 @@ TauNNId::TauNNId(const std::string &iInput, const TauNNTFCache *cache, const std
   fInput_ = iInput;
 }
 
-TauNNId::~TauNNId() { tensorflow::closeSession(session_); }
 void TauNNId::setNNVectorVar() {
   NNvectorVar_.clear();
   for (int i0 = 0; i0 < fNParticles_; i0++) {