forward-porting DNN-related from branch cms-tau-pog:CMSSW_9_4_X_tau_pog_DNNTauIDs

RecoTauTag/RecoTau/interface/DeepTauBase.h

@@ -14,6 +14,7 @@
 #include "FWCore/Framework/interface/stream/EDProducer.h"
 #include "FWCore/ParameterSet/interface/ParameterSet.h"
 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+#include "tensorflow/core/util/memmapped_file_system.h"
 #include "DataFormats/PatCandidates/interface/Electron.h"
 #include "DataFormats/PatCandidates/interface/Muon.h"
 #include "DataFormats/PatCandidates/interface/Tau.h"
@@ -22,10 +23,39 @@
 #include "RecoTauTag/RecoTau/interface/PFRecoTauClusterVariables.h"
 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
+#include <TF1.h>
 
 namespace deep_tau {
 
-class DeepTauBase : public edm::stream::EDProducer<> {
+class TauWPThreshold {
+public:
+    explicit TauWPThreshold(const std::string& cut_str);
+    double operator()(const pat::Tau& tau) const;
+
+private:
+    std::unique_ptr<TF1> fn_;
+    double value_;
+};
+
+class DeepTauCache {
+public:
+    using GraphPtr = std::shared_ptr<tensorflow::GraphDef>;
+
+    DeepTauCache(const std::string& graph_name, bool mem_mapped);
+    ~DeepTauCache();
+
+   // A Session allows concurrent calls to Run(), though a Session must
+   // be created / extended by a single thread.
+   tensorflow::Session& getSession() const { return *session_; }
+   const tensorflow::GraphDef& getGraph() const { return *graph_; }
+
+private:
+    GraphPtr graph_;
+    tensorflow::Session* session_;
+    std::unique_ptr<tensorflow::MemmappedEnv> memmappedEnv_;
+};
+
+class DeepTauBase : public edm::stream::EDProducer<edm::GlobalCache<DeepTauCache>> {
 public:
     using TauType = pat::Tau;
     using TauDiscriminator = pat::PATTauDiscriminator;
@@ -35,17 +65,15 @@ class DeepTauBase : public edm::stream::EDProducer<> {
     using ElectronCollection = pat::ElectronCollection;
     using MuonCollection = pat::MuonCollection;
     using LorentzVectorXYZ = ROOT::Math::LorentzVector<ROOT::Math::PxPyPzE4D<double>>;
-    using GraphPtr = std::shared_ptr<tensorflow::GraphDef>;
-    using Cutter = StringObjectFunction<TauType>;
+    using Cutter = TauWPThreshold;
     using CutterPtr = std::unique_ptr<Cutter>;
     using WPMap = std::map<std::string, CutterPtr>;
 
-
     struct Output {
         using ResultMap = std::map<std::string, std::unique_ptr<TauDiscriminator>>;
-        std::vector<size_t> num, den;
+        std::vector<size_t> num_, den_;
 
-        Output(const std::vector<size_t>& _num, const std::vector<size_t>& _den) : num(_num), den(_den) {}
+        Output(const std::vector<size_t>& num, const std::vector<size_t>& den) : num_(num), den_(den) {}
 
         ResultMap get_value(const edm::Handle<TauCollection>& taus, const tensorflow::Tensor& pred,
                             const WPMap& working_points) const;
@@ -54,27 +82,25 @@ class DeepTauBase : public edm::stream::EDProducer<> {
     using OutputCollection = std::map<std::string, Output>;
 
 
-    DeepTauBase(const edm::ParameterSet& cfg, const OutputCollection& outputs);
-    virtual ~DeepTauBase();
+    DeepTauBase(const edm::ParameterSet& cfg, const OutputCollection& outputs, const DeepTauCache* cache);
+    virtual ~DeepTauBase() {}
 
     virtual void produce(edm::Event& event, const edm::EventSetup& es) override;
 
+    static std::unique_ptr<DeepTauCache> initializeGlobalCache(const edm::ParameterSet& cfg);
+    static void globalEndJob(const DeepTauCache* cache){ }
 private:
-    virtual tensorflow::Tensor GetPredictions(edm::Event& event, const edm::EventSetup& es,
+    virtual tensorflow::Tensor getPredictions(edm::Event& event, const edm::EventSetup& es,
                                               edm::Handle<TauCollection> taus) = 0;
-    virtual void CreateOutputs(edm::Event& event, const tensorflow::Tensor& pred, edm::Handle<TauCollection> taus);
+    virtual void createOutputs(edm::Event& event, const tensorflow::Tensor& pred, edm::Handle<TauCollection> taus);
 
 protected:
-    edm::EDGetTokenT<TauCollection> taus_token;
-    std::string graphName;
-    GraphPtr graph;
-    tensorflow::Session* session;
-    std::map<std::string, WPMap> working_points;
-    OutputCollection outputs;
+    edm::EDGetTokenT<TauCollection> tausToken_;
+    std::map<std::string, WPMap> workingPoints_;
+    OutputCollection outputs_;
+    const DeepTauCache* cache_;
 };
 
 } // namespace deep_tau
 
-
-
 #endif