From 93abcd89b88b5f1cbc773088cc8ee8bf63f948f3 Mon Sep 17 00:00:00 2001 From: Joosep Pata Date: Fri, 6 Nov 2020 13:12:49 +0200 Subject: [PATCH] update Sonic producer, customization functions --- .../PFProducer/plugins/MLPFProducerSonic.cc | 130 ++++++++++-------- .../python/mlpfproducer_customise.py | 32 ++++- 2 files changed, 96 insertions(+), 66 deletions(-) diff --git a/RecoParticleFlow/PFProducer/plugins/MLPFProducerSonic.cc b/RecoParticleFlow/PFProducer/plugins/MLPFProducerSonic.cc index 91a6e759165d2..c8dd85c95791e 100644 --- a/RecoParticleFlow/PFProducer/plugins/MLPFProducerSonic.cc +++ b/RecoParticleFlow/PFProducer/plugins/MLPFProducerSonic.cc @@ -36,7 +36,21 @@ // git clone https://github.com/jpata/mlpf-models.git // mv mlpf-models/triton/models ./ // singularity exec --nv docker://nvcr.io/nvidia/tritonserver:20.10-py3 tritonserver --model-repository models --log-verbose=1 --log-error=1 --log-info=1 -// cmsRun src/RecoParticleFlow/PFProducer/test/mlpf_sonic.py +// +// process.mlpfproducergpu = cms.EDProducer("MLPFProducerSonicTriton", +// Client = cms.PSet( +// mode = cms.string("Sync"), +// batchSize = cms.untracked.uint32(1), +// address = cms.untracked.string("127.0.0.1"), +// port = cms.untracked.uint32(8001), +// timeout = cms.untracked.uint32(60), +// modelName = cms.string("mlpf"), +// modelVersion = cms.string("1"), +// verbose = cms.untracked.bool(False), +// allowedTries = cms.untracked.uint32(0), +// ) +// ) + class MLPFProducerSonicTriton : public SonicEDProducer { public: @@ -104,76 +118,70 @@ class MLPFProducerSonicTriton : public SonicEDProducer { ielem += 1; } input1.toServer(data1); - edm::LogInfo("MLPFProducerSonic") << "acquire done"; } void produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) override { - edm::LogInfo(client_.debugName()) << "produce"; + using namespace reco::mlpf; + std::vector pOutputCandidateCollection; const auto& output1 = iOutput.at("Identity"); - const auto& tmp = output1.template fromServer(); - for (const auto s : output1.shape()) { - edm::LogInfo(client_.debugName()) << "output shape=" << s; - } + + //get the data of the first (and only) batch + const auto& out_data = output1.template fromServer(); + const auto num_elem = output1.shape()[1]; - std::cout << "tmp.size" << tmp.size() << std::endl; - std::cout << "output.shape=" << output1.shape()[0] << "," << output1.shape()[1] << "," << output1.shape()[2] - << std::endl; - for (long int ielem = 0; ielem < num_elem; ielem++) { - for (long int iprop = 0; iprop < 12; iprop++) { - std::cout << tmp[0][ielem * output1.shape()[2] + iprop] << " "; + unsigned int num_pred_particles = 0; + for (unsigned int ielem = 0; ielem < num_elem; ielem++) { + //get the coefficients in the output corresponding to the class probabilities (raw logits) + std::vector pred_id_logits; + for (unsigned int idx_id = 0; idx_id <= NUM_CLASS; idx_id++) { + pred_id_logits.push_back(out_data[0][ielem*NUM_OUTPUTS + idx_id]); } - std::cout << std::endl; - } - //edm::LogInfo(client_.debugName()) << "produced shape=" << tmp.size(); - // edm::LogInfo(client_.debugName()) << "batchSize=" << client_.batchSize() << " noutput=" << client_.noutput(); - - // float output[client_.batchSize()][1000][NUM_OUTPUTS]; - - // for (unsigned int ibatch = 0; ibatch < client_.batchSize(); ibatch++) { - // for (unsigned int ielem = 0; ielem < 1000; ielem++) { - // for (unsigned int iprop = 0; iprop < NUM_OUTPUTS; iprop++) { - // output[ibatch][ielem][iprop] = iOutput.at(ibatch*1000 + ielem*NUM_OUTPUTS + iprop); - // } - // } - // } - // std::stringstream msg; - // for (unsigned int ibatch=0; ibatch pred_id_logits; - // for (unsigned int idx_id=0; idx_id <= NUM_CLASS; idx_id++) { - // pred_id_logits.push_back(output[ibatch][ielem][idx_id]); - // } - // //get the most probable class PDGID - // int pred_pid = pdgid_encoding.at(arg_max(pred_id_logits)); - // if (pred_pid != 0) { - // float pred_eta = output[ibatch][ielem][IDX_ETA]; - // float pred_phi = output[ibatch][ielem][IDX_PHI]; - // pred_phi = atan2(sin(pred_phi), cos(pred_phi)); - // float pred_e = output[ibatch][ielem][IDX_ENERGY]; - // float pred_charge = output[ibatch][ielem][IDX_CHARGE]; - // float pred_pt = pred_e / cosh(pred_eta); - // reco::PFCandidate::Charge charge = 0; - // if (pred_pid == 11 || pred_pid == 13 || pred_pid == 211) { - // charge = pred_charge > 0 ? +1 : -1; - // } - // TLorentzVector p4; - // p4.SetPtEtaPhiE(pred_pt, pred_eta, pred_phi, pred_e); - - // reco::PFCandidate cand(0, math::XYZTLorentzVector(p4.X(), p4.Y(), p4.Z(), p4.E()), reco::PFCandidate::ParticleType(0)); - // cand.setPdgId(pred_pid); - // cand.setCharge(charge); - // pOutputCandidateCollection.push_back(cand); - // } - // } - // } - // iEvent.emplace(pfCandidatesToken_, pOutputCandidateCollection); - - //edm::LogInfo(client_.debugName()) << msg.str(); - edm::LogInfo(client_.debugName()) << "produce done"; + //get the most probable class PDGID + int pred_pid = pdgid_encoding[arg_max(pred_id_logits)]; + + //get the predicted momentum components + float pred_eta = out_data[0][ielem*NUM_OUTPUTS + IDX_ETA]; + float pred_phi = out_data[0][ielem*NUM_OUTPUTS + IDX_PHI]; + + //put phi into a modular range + pred_phi = TMath::ATan2(TMath::Sin(pred_phi), TMath::Cos(pred_phi)); + float pred_e = out_data[0][ielem*NUM_OUTPUTS + IDX_ENERGY]; + + //currently, set the pT from a massless approximation. + //later versions of the model may predict predict both the energy and pT of the particle + float pred_pt = pred_e / cosh(pred_eta); + float pred_charge = out_data[0][ielem*NUM_OUTPUTS + IDX_CHARGE]; + + //a particle was predicted for this PFElement + if (pred_pid != 0) { + //set the charge to +1 or -1 for PFCandidates that are charged, according to the sign of the predicted charge + reco::PFCandidate::Charge charge = 0; + if (pred_pid == 11 || pred_pid == 13 || pred_pid == 211) { + charge = pred_charge > 0 ? +1 : -1; + } + + TLorentzVector p4; + p4.SetPtEtaPhiE(pred_pt, pred_eta, pred_phi, pred_e); + + reco::PFCandidate cand( + 0, math::XYZTLorentzVector(p4.X(), p4.Y(), p4.Z(), p4.E()), reco::PFCandidate::ParticleType(0)); + cand.setPdgId(pred_pid); + cand.setCharge(charge); + + pOutputCandidateCollection.push_back(cand); + num_pred_particles += 1; + } else { + //this element did not directly yield a particle, but may have been used indirectly for other ML-PFCandidates + //we can determine this from the sparse adjacency matrix that the model produces internally + } + } //loop over PFElements + + iEvent.emplace(pfCandidatesToken_, pOutputCandidateCollection); } + ~MLPFProducerSonicTriton() override {} //to ensure distinct cfi names - specialized below diff --git a/RecoParticleFlow/PFProducer/python/mlpfproducer_customise.py b/RecoParticleFlow/PFProducer/python/mlpfproducer_customise.py index 82aba41c842f0..c9bfab4449757 100644 --- a/RecoParticleFlow/PFProducer/python/mlpfproducer_customise.py +++ b/RecoParticleFlow/PFProducer/python/mlpfproducer_customise.py @@ -1,6 +1,6 @@ import FWCore.ParameterSet.Config as cms -def customise_step3(process, output="AODSIMoutput"): +def customise_step3(process, output="AODSIMoutput", usegpu=False): if hasattr(process, "particleFlowBlock") and hasattr(process, output): process.mlpfproducer = cms.EDProducer("MLPFProducer", #we don't really need anything the block does, just to get all the PFElements in the event in a convenient list @@ -16,11 +16,33 @@ def customise_step3(process, output="AODSIMoutput"): out.outputCommands.append('keep *_ak4PFJets_*_*') process.mlpf_path = cms.Path(process.mlpfproducer*process.ak4MLPFJets) + if usegpu: + process.mlpfproducergpu = cms.EDProducer("MLPFProducerSonicTriton", + Client = cms.PSet( + mode = cms.string("Sync"), + batchSize = cms.untracked.uint32(1), + address = cms.untracked.string("127.0.0.1"), + port = cms.untracked.uint32(8001), + timeout = cms.untracked.uint32(60), + modelName = cms.string("mlpf"), + modelVersion = cms.string("1"), + verbose = cms.untracked.bool(False), + allowedTries = cms.untracked.uint32(0), + ) + ) + process.mlpf_path.insert(-1, process.mlpfproducergpu) + process.schedule.insert(-1, process.mlpf_path) return process -def customise_step3_reco(process, output="RECOSIMoutput"): - return customise_step3(process, output) +def customise_step3_reco(process): + return customise_step3(process, "RECOSIMoutput") + +def customise_step3_aod(process): + return customise_step3(process, "AODSIMoutput") + +def customise_step3_reco_gpu(process): + return customise_step3(process, "RECOSIMoutput", True) -def customise_step3_aod(process, output="AODSIMoutput"): - return customise_step3(process, output) +def customise_step3_aod_gpu(process): + return customise_step3(process, "AODSIMoutput", True) \ No newline at end of file