Skip to content

Commit

Permalink
Merge pull request cms-sw#105 from MRD2F/CMSSW_9_4_X_DPFIso_deepTau
Browse files Browse the repository at this point in the history 
Improvent of memory usage
  • Loading branch information
@mbluj
mbluj authored Nov 13, 2018
2 parents 42d930e + 80c745a commit 3044c0f
Show file tree
Hide file tree
Showing 6 changed files with 355 additions and 196 deletions.
62 changes: 44 additions & 18 deletions RecoTauTag/RecoTau/interface/DeepTauBase.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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, const 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_; }

protected:
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;
Expand All @@ -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;
Expand All @@ -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
56 changes: 23 additions & 33 deletions RecoTauTag/RecoTau/plugins/DPFIsolation.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,29 +12,24 @@ namespace {
inline int getPFCandidateIndex(const edm::Handle<pat::PackedCandidateCollection>& pfcands,
const reco::CandidatePtr& cptr)
{
unsigned int pfInd = -1;
for(unsigned int i = 0; i < pfcands->size(); ++i) {
pfInd++;
if(reco::CandidatePtr(pfcands,i) == cptr) {
pfInd = i;
break;
}
if(reco::CandidatePtr(pfcands,i) == cptr)
return i;
}
return pfInd;
return -1;
}
} // anonymous namespace


class DPFIsolation : public deep_tau::DeepTauBase {
public:
static OutputCollection& GetOutputs()
{
static size_t tau_index = 0;
static OutputCollection outputs = { { "VSall", Output({tau_index}, {}) } };
return outputs;
static OutputCollection outputs_ = { { "VSall", Output({tau_index}, {}) } };
return outputs_;
};

static unsigned GetNumberOfParticles(unsigned graphVersion)
static unsigned getNumberOfParticles(unsigned graphVersion)
{
static const std::map<unsigned, unsigned> nparticles { { 0, 60 }, { 1, 36 } };
return nparticles.at(graphVersion);
Expand All @@ -52,7 +47,9 @@ class DPFIsolation : public deep_tau::DeepTauBase {
desc.add<edm::InputTag>("pfcands", edm::InputTag("packedPFCandidates"));
desc.add<edm::InputTag>("taus", edm::InputTag("slimmedTaus"));
desc.add<edm::InputTag>("vertices", edm::InputTag("offlineSlimmedPrimaryVertices"));
desc.add<std::string>("graph_file", "RecoTauTag/TrainingFiles/data/DPFTauId/DPFIsolation_2017v0.pb");
desc.add<std::string>("graph_file", "RecoTauTag/TrainingFiles/data/DPFTauId/DPFIsolation_2017v0_quantized.pb");
desc.add<unsigned>("version", 0);
desc.add<bool>("mem_mapped", false);

edm::ParameterSetDescription descWP;
descWP.add<std::string>("VVVLoose", "0");
Expand All @@ -68,21 +65,18 @@ class DPFIsolation : public deep_tau::DeepTauBase {
descriptions.add("DPFTau2016v0", desc);
}

explicit DPFIsolation(const edm::ParameterSet& cfg) :
DeepTauBase(cfg, GetOutputs()),
explicit DPFIsolation(const edm::ParameterSet& cfg,const deep_tau::DeepTauCache* cache) :
DeepTauBase(cfg, GetOutputs(), cache),
pfcand_token(consumes<pat::PackedCandidateCollection>(cfg.getParameter<edm::InputTag>("pfcands"))),
vtx_token(consumes<reco::VertexCollection>(cfg.getParameter<edm::InputTag>("vertices")))
vtx_token(consumes<reco::VertexCollection>(cfg.getParameter<edm::InputTag>("vertices"))),
graphVersion(cfg.getParameter<unsigned>("version"))
{
if(graphName.find("v0.pb") != std::string::npos)
graphVersion = 0;
else if(graphName.find("v1.pb") != std::string::npos)
graphVersion = 1;
else
throw cms::Exception("DPFIsolation") << "unknown version of the graph file.";
if(!(graphVersion == 1 || graphVersion == 0 ))
throw cms::Exception("DPFIsolation") << "unknown version of the graph_ file.";
}

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) override
{
edm::Handle<pat::PackedCandidateCollection> pfcands;
Expand All @@ -92,11 +86,11 @@ class DPFIsolation : public deep_tau::DeepTauBase {
event.getByToken(vtx_token, vertices);

tensorflow::Tensor tensor(tensorflow::DT_FLOAT, {1,
static_cast<int>(GetNumberOfParticles(graphVersion)), static_cast<int>(GetNumberOfFeatures(graphVersion))});
static_cast<int>(getNumberOfParticles(graphVersion)), static_cast<int>(GetNumberOfFeatures(graphVersion))});

tensorflow::Tensor predictions(tensorflow::DT_FLOAT, { static_cast<int>(taus->size()), 1});

std::vector<tensorflow::Tensor> outputs;
std::vector<tensorflow::Tensor> outputs_;

float pfCandPt, pfCandPz, pfCandPtRel, pfCandPzRel, pfCandDr, pfCandDEta, pfCandDPhi, pfCandEta, pfCandDz,
pfCandDzErr, pfCandD0, pfCandD0D0, pfCandD0Dz, pfCandD0Dphi, pfCandPuppiWeight,
Expand Down Expand Up @@ -124,20 +118,20 @@ class DPFIsolation : public deep_tau::DeepTauBase {

std::vector<unsigned int> signalCandidateInds;

for(auto c : tau.signalCands())
for(const auto c : tau.signalCands())
signalCandidateInds.push_back(getPFCandidateIndex(pfcands,c));

float lepRecoPt = tau.pt();
float lepRecoPz = std::abs(tau.pz());

// Use of setZero results in warnings in eigen library during compilation.
//tensor.flat<float>().setZero();
const unsigned n_inputs = GetNumberOfParticles(graphVersion) * GetNumberOfFeatures(graphVersion);
const unsigned n_inputs = getNumberOfParticles(graphVersion) * GetNumberOfFeatures(graphVersion);
for(unsigned input_idx = 0; input_idx < n_inputs; ++input_idx)
tensor.flat<float>()(input_idx) = 0;

unsigned int iPF = 0;
const unsigned max_iPF = GetNumberOfParticles(graphVersion);
const unsigned max_iPF = getNumberOfParticles(graphVersion);

std::vector<unsigned int> sorted_inds(pfcands->size());
std::size_t n = 0;
Expand All @@ -153,8 +147,6 @@ class DPFIsolation : public deep_tau::DeepTauBase {
if (p.pt() < 0.5) continue;
if (p.fromPV() < 0) continue;
if (deltaR_tau_p > 0.5) continue;


if (p.fromPV() < 1 && p.charge() != 0) continue;
pfCandPt = p.pt();
pfCandPtRel = p.pt()/lepRecoPt;
Expand Down Expand Up @@ -372,12 +364,10 @@ class DPFIsolation : public deep_tau::DeepTauBase {
tensor.tensor<float,3>()( 0, 36-1-iPF, 49) = pfCandPdgID==211;
tensor.tensor<float,3>()( 0, 36-1-iPF, 50) = pfCandTauIndMatch;
}

iPF++;
}

tensorflow::Status status = session->Run( { {"input_1", tensor} }, {"output_node0"}, {}, &outputs);
predictions.matrix<float>()(tau_index, 0) = outputs[0].flat<float>()(0);
tensorflow::run(&(cache_->getSession()), { { "input_1", tensor } }, { "output_node0" }, {}, &outputs_);
predictions.matrix<float>()(tau_index, 0) = outputs_[0].flat<float>()(0);
}
return predictions;
}
Expand Down
Loading

0 comments on commit 3044c0f

Please sign in to comment.