diff --git a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py index 396fcefd84a79..8f8df61d5be6d 100644 --- a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py +++ b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py @@ -5,6 +5,9 @@ from Configuration.ProcessModifiers.particleNetPTSonicTriton_cff import particleNetPTSonicTriton from Configuration.ProcessModifiers.deepMETSonicTriton_cff import deepMETSonicTriton from Configuration.ProcessModifiers.deepTauSonicTriton_cff import deepTauSonicTriton +from Configuration.ProcessModifiers.particleTransformerAK4SonicTriton_cff import particleTransformerAK4SonicTriton +from Configuration.ProcessModifiers.unifiedparticleTransformerAK4SonicTriton_cff import unifiedparticleTransformerAK4SonicTriton + # collect all SonicTriton-related process modifiers here -allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,deepTauSonicTriton) +allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton,unifiedparticleTransformerAK4SonicTriton) diff --git a/Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py b/Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py new file mode 100644 index 0000000000000..22fff0fb9ea6c --- /dev/null +++ b/Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py @@ -0,0 +1,3 @@ +import FWCore.ParameterSet.Config as cms + +particleTransformerAK4SonicTriton = cms.Modifier() diff --git a/Configuration/ProcessModifiers/python/unifiedparticleTransformerAK4SonicTriton_cff.py b/Configuration/ProcessModifiers/python/unifiedparticleTransformerAK4SonicTriton_cff.py new file mode 100644 index 0000000000000..d7b082a010613 --- /dev/null +++ b/Configuration/ProcessModifiers/python/unifiedparticleTransformerAK4SonicTriton_cff.py @@ -0,0 +1,3 @@ +import FWCore.ParameterSet.Config as cms + +unifiedparticleTransformerAK4SonicTriton = cms.Modifier() diff --git a/RecoBTag/ONNXRuntime/BuildFile.xml b/RecoBTag/ONNXRuntime/BuildFile.xml index c9b2daa7fe264..45e88927184fa 100644 --- a/RecoBTag/ONNXRuntime/BuildFile.xml +++ b/RecoBTag/ONNXRuntime/BuildFile.xml @@ -1,3 +1,4 @@ + diff --git a/RecoBTag/ONNXRuntime/interface/tensor_configs.h b/RecoBTag/ONNXRuntime/interface/tensor_configs.h index e528b4e92f74a..c5c14006cbbf2 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_configs.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_configs.h @@ -1,6 +1,8 @@ #ifndef RecoBTag_ONNXRuntime_tensor_configs_h #define RecoBTag_ONNXRuntime_tensor_configs_h +#include + namespace deepflavour { constexpr unsigned n_features_global = 15; @@ -28,4 +30,85 @@ namespace deepvertex { } // namespace deepvertex +namespace parT { + + enum InputFeatures { + kBegin=0, + kChargedCandidates=kBegin, + kNeutralCandidates=1, + kVertices=2, + kChargedCandidates4Vec=3, + kNeutralCandidates4Vec=4, + kVertices4Vec=5, + kEnd=6 + }; + + constexpr unsigned n_cpf_accept = 25; + constexpr unsigned n_npf_accept = 25; + constexpr unsigned n_sv_accept = 5; + + const std::map N_InputFeatures{ + {kChargedCandidates, 16}, + {kNeutralCandidates, 8}, + {kVertices, 14}, + {kChargedCandidates4Vec, 4}, + {kNeutralCandidates4Vec, 4}, + {kVertices4Vec, 4} + }; + + const std::map N_AcceptedFeatures{ + {kChargedCandidates, n_cpf_accept}, + {kNeutralCandidates, n_npf_accept}, + {kVertices, n_sv_accept}, + {kChargedCandidates4Vec, n_cpf_accept}, + {kNeutralCandidates4Vec, n_npf_accept}, + {kVertices4Vec, n_sv_accept} + }; + +} // namespace parT + + +namespace UparT { + + enum InputFeatures { + kBegin=0, + kChargedCandidates=kBegin, + kLostTracks = 1, + kNeutralCandidates=2, + kVertices=3, + kChargedCandidates4Vec=4, + kLostTracks4Vec=5, + kNeutralCandidates4Vec=6, + kVertices4Vec=7, + kEnd=8 + }; + + constexpr unsigned n_cpf_accept = 29; + constexpr unsigned n_lt_accept = 5; + constexpr unsigned n_npf_accept = 25; + constexpr unsigned n_sv_accept = 5; + + const std::map N_InputFeatures{ + {kChargedCandidates, 25}, + {kLostTracks, 18}, + {kNeutralCandidates, 8}, + {kVertices, 14}, + {kChargedCandidates4Vec, 4}, + {kLostTracks4Vec, 4}, + {kNeutralCandidates4Vec, 4}, + {kVertices4Vec, 4} + }; + + const std::map N_AcceptedFeatures{ + {kChargedCandidates, n_cpf_accept}, + {kLostTracks, n_lt_accept}, + {kNeutralCandidates, n_npf_accept}, + {kVertices, n_sv_accept}, + {kChargedCandidates4Vec, n_cpf_accept}, + {kLostTracks4Vec, n_lt_accept}, + {kNeutralCandidates4Vec, n_npf_accept}, + {kVertices4Vec, n_sv_accept} + }; + +} // namespace UparT #endif diff --git a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h index bb1a25a2d0f0b..62100573bcd5b 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h @@ -2,6 +2,10 @@ #define RecoBTag_ONNXRuntime_tensor_fillers_h #include "DataFormats/BTauReco/interface/DeepFlavourTagInfo.h" +#include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h" +#include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4Features.h" +#include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h" +#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h" namespace btagbtvdeep { @@ -19,6 +23,79 @@ namespace btagbtvdeep { void neighbourTrack_tensor_filler(float*& ptr, const btagbtvdeep::TrackPairFeatures& neighbourTrack_features); + std::vector inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputFeatures ifeature); + + std::vector inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputFeatures ifeature); + + std::vector inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature); + + std::vector inputs_UparT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, UparT::InputFeatures ifeature); + + std::vector inputs_UparT(const btagbtvdeep::LostTracksFeatures& lt_features, UparT::InputFeatures ifeature); + + std::vector inputs_UparT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, UparT::InputFeatures ifeature); + + std::vector inputs_UparT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, UparT::InputFeatures ifeature); + + template + void parT_tensor_filler(cms::Ort::FloatArrays& data, const parT::InputFeatures ifeature, const std::vector& features, const unsigned int max_n, const float*& start, unsigned offset) { + float* ptr = nullptr; + for (std::size_t n = 0; n < max_n; n++) { + const auto& f = features.at(n); + ptr = &data[ifeature][offset + n * parT::N_InputFeatures.at(ifeature)]; + start = ptr; + const std::vector& inputs = inputs_parT(f, ifeature); + for (unsigned int i = 0; i < inputs.size(); i++) { + *ptr = inputs[i]; + ++ptr; + } + if (inputs.size() > 0) --ptr; + assert(start + parT::N_InputFeatures.at(ifeature) - 1 == ptr); + } + } + + template + void parT_tensor_filler(std::vector& vdata, const parT::InputFeatures ifeature, const std::vector& features, const unsigned int target_n) { + unsigned int n = std::clamp((unsigned int)features.size(), (unsigned int)0, (unsigned int)parT::N_AcceptedFeatures.at(ifeature)); + for (unsigned int count = 0; count < n; count++) { + const std::vector& inputs = inputs_parT(features.at(count), ifeature); + vdata.insert(vdata.end(), inputs.begin(), inputs.end()); + } + unsigned int n_features = parT::N_InputFeatures.at(ifeature); + if (n < target_n) + vdata.insert(vdata.end(), (target_n - n) * n_features, 0); // Add 0 to unfilled part as padding value + } + + template + void UparT_tensor_filler(cms::Ort::FloatArrays& data, const UparT::InputFeatures ifeature, const std::vector& features, const unsigned int max_n, const float*& start, unsigned offset) { + float* ptr = nullptr; + for (std::size_t n = 0; n < max_n; n++) { + const auto& f = features.at(n); + ptr = &data[ifeature][offset + n * UparT::N_InputFeatures.at(ifeature)]; + start = ptr; + const std::vector& inputs = inputs_UparT(f, ifeature); + for (unsigned int i = 0; i < inputs.size(); i++) { + *ptr = inputs[i]; + ++ptr; + } + if (inputs.size() > 0) --ptr; + assert(start + UparT::N_InputFeatures.at(ifeature) - 1 == ptr); + } + } + + template + void UparT_tensor_filler(std::vector& vdata, const UparT::InputFeatures ifeature, const std::vector& features, const unsigned int target_n) { + unsigned int n = std::clamp((unsigned int)features.size(), (unsigned int)0, (unsigned int)UparT::N_AcceptedFeatures.at(ifeature)); + for (unsigned int count = 0; count < n; count++) { + const std::vector& inputs = inputs_UparT(features.at(count), ifeature); + vdata.insert(vdata.end(), inputs.begin(), inputs.end()); + } + unsigned int n_features = UparT::N_InputFeatures.at(ifeature); + if (n < target_n) + vdata.insert(vdata.end(), (target_n - n) * n_features, 0); // Add 0 to unfilled part as padding value + } + + } // namespace btagbtvdeep #endif diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc index b27a9b397e520..e63b5504936b4 100644 --- a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc +++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc @@ -16,6 +16,9 @@ #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h" +#include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h" +#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h" + using namespace cms::Ort; class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer> { @@ -27,7 +30,7 @@ class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer static std::unique_ptr initializeGlobalCache(const edm::ParameterSet&); static void globalEndJob(const ONNXRuntime*); - + private: typedef std::vector TagInfoCollection; typedef reco::JetTagCollection JetTagCollection; @@ -41,24 +44,9 @@ class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer std::vector flav_names_; std::vector input_names_; std::vector output_names_; - - enum InputIndexes { - kChargedCandidates = 0, - kNeutralCandidates = 1, - kVertices = 2, - kChargedCandidates4Vec = 3, - kNeutralCandidates4Vec = 4, - kVertices4Vec = 5 - }; - unsigned n_cpf_; - constexpr static unsigned n_features_cpf_ = 16; - constexpr static unsigned n_pairwise_features_cpf_ = 4; - unsigned n_npf_; - constexpr static unsigned n_features_npf_ = 8; - constexpr static unsigned n_pairwise_features_npf_ = 4; - unsigned n_sv_; - constexpr static unsigned n_features_sv_ = 14; - constexpr static unsigned n_pairwise_features_sv_ = 4; + unsigned int n_cpf_; + unsigned int n_npf_; + unsigned int n_sv_; std::vector input_sizes_; std::vector> input_shapes_; // shapes of each input group (-1 for dynamic axis) @@ -84,7 +72,7 @@ void ParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::Configurat desc.add("src", edm::InputTag("pfParticleTransformerAK4TagInfos")); desc.add>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"}); desc.add("model_path", - edm::FileInPath("RecoBTag/Combined/data/RobustParTAK4/PUPPI/V00/RobustParTAK4.onnx")); + edm::FileInPath("RecoBTag/Combined/data/RobustParTAK4/PUPPI/V00/modelfile/model.onnx")); desc.add>("output_names", {"softmax"}); desc.add>( "flav_names", std::vector{"probb", "probbb", "problepb", "probc", "probuds", "probg"}); @@ -124,12 +112,12 @@ void ParticleTransformerAK4ONNXJetTagsProducer::produce(edm::Event& iEvent, cons get_input_sizes(taginfo); // run prediction with dynamic batch size per event - input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)n_features_cpf_}, - {(int64_t)1, (int64_t)n_npf_, (int64_t)n_features_npf_}, - {(int64_t)1, (int64_t)n_sv_, (int64_t)n_features_sv_}, - {(int64_t)1, (int64_t)n_cpf_, (int64_t)n_pairwise_features_cpf_}, - {(int64_t)1, (int64_t)n_npf_, (int64_t)n_pairwise_features_npf_}, - {(int64_t)1, (int64_t)n_sv_, (int64_t)n_pairwise_features_sv_}}; + input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::N_InputFeatures.at(parT::kChargedCandidates)}, + {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::N_InputFeatures.at(parT::kNeutralCandidates)}, + {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::N_InputFeatures.at(parT::kVertices)}, + {(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::N_InputFeatures.at(parT::kChargedCandidates4Vec)}, + {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::N_InputFeatures.at(parT::kNeutralCandidates4Vec)}, + {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::N_InputFeatures.at(parT::kVertices4Vec)}}; outputs = globalCache()->run(input_names_, data_, input_shapes_, output_names_, 1)[0]; assert(outputs.size() == flav_names_.size()); @@ -151,24 +139,17 @@ void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes( const reco::FeaturesTagInfo taginfo) { const auto& features = taginfo.features(); - unsigned int n_cpf = features.c_pf_features.size(); - unsigned int n_npf = features.n_pf_features.size(); - unsigned int n_vtx = features.sv_features.size(); - - n_cpf_ = std::max((unsigned int)1, n_cpf); - n_npf_ = std::max((unsigned int)1, n_npf); - n_sv_ = std::max((unsigned int)1, n_vtx); + n_cpf_ = std::clamp((unsigned int)features.c_pf_features.size(), (unsigned int)1, (unsigned int)parT::n_cpf_accept); + n_npf_ = std::clamp((unsigned int)features.n_pf_features.size(), (unsigned int)1, (unsigned int)parT::n_npf_accept); + n_sv_ = std::clamp((unsigned int)features.sv_features.size(), (unsigned int)1, (unsigned int)parT::n_sv_accept); - n_cpf_ = std::min((unsigned int)25, n_cpf_); - n_npf_ = std::min((unsigned int)25, n_npf_); - n_sv_ = std::min((unsigned int)5, n_sv_); input_sizes_ = { - n_cpf_ * n_features_cpf_, - n_npf_ * n_features_npf_, - n_sv_ * n_features_sv_, - n_cpf_ * n_pairwise_features_cpf_, - n_npf_ * n_pairwise_features_npf_, - n_sv_ * n_pairwise_features_sv_, + n_cpf_ * parT::N_InputFeatures.at(parT::kChargedCandidates), + n_npf_ * parT::N_InputFeatures.at(parT::kNeutralCandidates), + n_sv_ * parT::N_InputFeatures.at(parT::kVertices), + n_cpf_ * parT::N_InputFeatures.at(parT::kChargedCandidates4Vec), + n_npf_ * parT::N_InputFeatures.at(parT::kNeutralCandidates4Vec), + n_sv_ * parT::N_InputFeatures.at(parT::kVertices4Vec), }; // init data storage data_.clear(); @@ -180,116 +161,25 @@ void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes( } void ParticleTransformerAK4ONNXJetTagsProducer::make_inputs(btagbtvdeep::ParticleTransformerAK4Features features) { - float* ptr = nullptr; const float* start = nullptr; unsigned offset = 0; - // c_pf candidates auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_); - for (std::size_t c_pf_n = 0; c_pf_n < max_c_pf_n; c_pf_n++) { - const auto& c_pf_features = features.c_pf_features.at(c_pf_n); - ptr = &data_[kChargedCandidates][offset + c_pf_n * n_features_cpf_]; - start = ptr; - *ptr = c_pf_features.btagPf_trackEtaRel; - *(++ptr) = c_pf_features.btagPf_trackPtRel; - *(++ptr) = c_pf_features.btagPf_trackPPar; - *(++ptr) = c_pf_features.btagPf_trackDeltaR; - *(++ptr) = c_pf_features.btagPf_trackPParRatio; - *(++ptr) = c_pf_features.btagPf_trackSip2dVal; - *(++ptr) = c_pf_features.btagPf_trackSip2dSig; - *(++ptr) = c_pf_features.btagPf_trackSip3dVal; - *(++ptr) = c_pf_features.btagPf_trackSip3dSig; - *(++ptr) = c_pf_features.btagPf_trackJetDistVal; - *(++ptr) = c_pf_features.ptrel; - *(++ptr) = c_pf_features.drminsv; - *(++ptr) = c_pf_features.vtx_ass; - *(++ptr) = c_pf_features.puppiw; - *(++ptr) = c_pf_features.chi2; - *(++ptr) = c_pf_features.quality; - assert(start + n_features_cpf_ - 1 == ptr); - } - - // n_pf candidates auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_); - for (std::size_t n_pf_n = 0; n_pf_n < max_n_pf_n; n_pf_n++) { - const auto& n_pf_features = features.n_pf_features.at(n_pf_n); - ptr = &data_[kNeutralCandidates][offset + n_pf_n * n_features_npf_]; - start = ptr; - *ptr = n_pf_features.ptrel; - *(++ptr) = n_pf_features.etarel; - *(++ptr) = n_pf_features.phirel; - *(++ptr) = n_pf_features.deltaR; - *(++ptr) = n_pf_features.isGamma; - *(++ptr) = n_pf_features.hadFrac; - *(++ptr) = n_pf_features.drminsv; - *(++ptr) = n_pf_features.puppiw; - assert(start + n_features_npf_ - 1 == ptr); - } - - // sv candidates auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_); - for (std::size_t sv_n = 0; sv_n < max_sv_n; sv_n++) { - const auto& sv_features = features.sv_features.at(sv_n); - ptr = &data_[kVertices][offset + sv_n * n_features_sv_]; - start = ptr; - *ptr = sv_features.pt; - *(++ptr) = sv_features.deltaR; - *(++ptr) = sv_features.mass; - *(++ptr) = sv_features.etarel; - *(++ptr) = sv_features.phirel; - *(++ptr) = sv_features.ntracks; - *(++ptr) = sv_features.chi2; - *(++ptr) = sv_features.normchi2; - *(++ptr) = sv_features.dxy; - *(++ptr) = sv_features.dxysig; - *(++ptr) = sv_features.d3d; - *(++ptr) = sv_features.d3dsig; - *(++ptr) = sv_features.costhetasvpv; - *(++ptr) = sv_features.enratio; - assert(start + n_features_sv_ - 1 == ptr); - } + // c_pf candidates + parT_tensor_filler(data_, parT::kChargedCandidates, features.c_pf_features, max_c_pf_n, start, offset); + // n_pf candidates + parT_tensor_filler(data_, parT::kNeutralCandidates, features.n_pf_features, max_n_pf_n, start, offset); + // sv candidates + parT_tensor_filler(data_, parT::kVertices, features.sv_features, max_sv_n, start, offset); // cpf pairwise features (4-vectors) - auto max_cpf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_); - for (std::size_t cpf_n = 0; cpf_n < max_cpf_n; cpf_n++) { - const auto& cpf_pairwise_features = features.c_pf_features.at(cpf_n); - ptr = &data_[kChargedCandidates4Vec][offset + cpf_n * n_pairwise_features_cpf_]; - start = ptr; - *ptr = cpf_pairwise_features.px; - *(++ptr) = cpf_pairwise_features.py; - *(++ptr) = cpf_pairwise_features.pz; - *(++ptr) = cpf_pairwise_features.e; - - assert(start + n_pairwise_features_cpf_ - 1 == ptr); - } - + parT_tensor_filler(data_, parT::kChargedCandidates4Vec, features.c_pf_features, max_c_pf_n, start, offset); // npf pairwise features (4-vectors) - auto max_npf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_); - for (std::size_t npf_n = 0; npf_n < max_npf_n; npf_n++) { - const auto& npf_pairwise_features = features.n_pf_features.at(npf_n); - ptr = &data_[kNeutralCandidates4Vec][offset + npf_n * n_pairwise_features_npf_]; - start = ptr; - *ptr = npf_pairwise_features.px; - *(++ptr) = npf_pairwise_features.py; - *(++ptr) = npf_pairwise_features.pz; - *(++ptr) = npf_pairwise_features.e; - - assert(start + n_pairwise_features_npf_ - 1 == ptr); - } - + parT_tensor_filler(data_, parT::kNeutralCandidates4Vec, features.n_pf_features, max_n_pf_n, start, offset); // sv pairwise features (4-vectors) - auto max_sv_N = std::min(features.sv_features.size(), (std::size_t)n_sv_); - for (std::size_t sv_N = 0; sv_N < max_sv_N; sv_N++) { - const auto& sv_pairwise_features = features.sv_features.at(sv_N); - ptr = &data_[kVertices4Vec][offset + sv_N * n_pairwise_features_sv_]; - start = ptr; - *ptr = sv_pairwise_features.px; - *(++ptr) = sv_pairwise_features.py; - *(++ptr) = sv_pairwise_features.pz; - *(++ptr) = sv_pairwise_features.e; - - assert(start + n_pairwise_features_sv_ - 1 == ptr); - } + parT_tensor_filler(data_, parT::kVertices4Vec, features.sv_features, max_sv_n, start, offset); } //define this as a plug-in diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc new file mode 100644 index 0000000000000..cf515d170373e --- /dev/null +++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc @@ -0,0 +1,187 @@ +#include "FWCore/Framework/interface/Frameworkfwd.h" + +#include "FWCore/Framework/interface/Event.h" +#include "FWCore/Framework/interface/MakerMacros.h" + +#include "FWCore/Framework/interface/makeRefToBaseProdFrom.h" + +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/Utilities/interface/StreamID.h" + +#include "DataFormats/BTauReco/interface/JetTag.h" + +#include "DataFormats/BTauReco/interface/ParticleTransformerAK4TagInfo.h" +#include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h" + +#include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonData.h" + +#include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h" +#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h" + +class ParticleTransformerAK4SonicJetTagsProducer : public TritonEDProducer<> { +public: + explicit ParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet&); + ~ParticleTransformerAK4SonicJetTagsProducer() override; + + void acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) override; + + void produce(edm::Event &iEvent, edm::EventSetup const &iSetup, Output const &iOutput) override; + static void fillDescriptions(edm::ConfigurationDescriptions&); + +private: + typedef std::vector TagInfoCollection; + typedef reco::JetTagCollection JetTagCollection; + + const edm::EDGetTokenT src_; + std::vector flav_names_; + std::vector input_names_; + std::vector output_names_; + + bool skippedInference_ = false; +}; + +ParticleTransformerAK4SonicJetTagsProducer::ParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet& iConfig) + : TritonEDProducer<>(iConfig), + src_(consumes(iConfig.getParameter("src"))), + flav_names_(iConfig.getParameter>("flav_names")), + input_names_(iConfig.getParameter>("input_names")), + output_names_(iConfig.getParameter>("output_names")) { + // get output names from flav_names + for (const auto& flav_name : flav_names_) { + produces(flav_name); + } +} + +ParticleTransformerAK4SonicJetTagsProducer::~ParticleTransformerAK4SonicJetTagsProducer() {} + +void ParticleTransformerAK4SonicJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) { + // pfParticleTransformerAK4JetTags + edm::ParameterSetDescription desc; + TritonClient::fillPSetDescription(desc); + desc.add("src", edm::InputTag("pfParticleTransformerAK4TagInfos")); + desc.add>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"}); + desc.add>("output_names", {"softmax"}); + desc.add>( + "flav_names", std::vector{"probb", "probbb", "problepb", "probc", "probuds", "probg"}); + + descriptions.add("pfParticleTransformerAK4SonicJetTags", desc); +} + +void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) { + edm::Handle tag_infos; + iEvent.getByToken(src_, tag_infos); + client_->setBatchSize(tag_infos->size()); + skippedInference_ = false; + if (tag_infos->empty()) return; + + // Find the max n_cpf, n_npf and n_vtx among all the jets in an event. + unsigned int max_n_cpf_counter = 0; + unsigned int max_n_npf_counter = 0; + unsigned int max_n_vtx_counter = 0; + for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) { + max_n_cpf_counter = std::max(max_n_cpf_counter, + static_cast(((*tag_infos)[jet_n]).features().c_pf_features.size())); + max_n_npf_counter = std::max(max_n_npf_counter, + static_cast(((*tag_infos)[jet_n]).features().n_pf_features.size())); + max_n_vtx_counter = std::max(max_n_vtx_counter, + static_cast(((*tag_infos)[jet_n]).features().sv_features.size())); + } + + // If an event has no jet, or all jets has zero n_cpf, n_npf and n_vtx, the inference is skipped. + if (max_n_cpf_counter == 0 && max_n_npf_counter == 0 && max_n_vtx_counter == 0) { + client_->setBatchSize(0); + skippedInference_ = true; + return; + } + + // all the jets in the same event will fill up the same amount of n_cpf, n_npf, n_vtx and send to server + const unsigned int target_n_cpf = std::clamp(max_n_cpf_counter, (unsigned int)1, (unsigned int)parT::n_cpf_accept); + const unsigned int target_n_npf = std::clamp(max_n_npf_counter, (unsigned int)1, (unsigned int)parT::n_npf_accept); + const unsigned int target_n_vtx = std::clamp(max_n_vtx_counter, (unsigned int)1, (unsigned int)parT::n_sv_accept); + + const std::map target_n { + {parT::kChargedCandidates, target_n_cpf}, + {parT::kNeutralCandidates, target_n_npf}, + {parT::kVertices, target_n_vtx}, + {parT::kChargedCandidates4Vec, target_n_cpf}, + {parT::kNeutralCandidates4Vec, target_n_npf}, + {parT::kVertices4Vec, target_n_vtx} + }; + + // loop through all groups of features + for (parT::InputFeatures ifeature = parT::kBegin; ifeature != parT::kEnd; ifeature = static_cast(ifeature+1)) { + const auto &group_name = input_names_[ifeature]; + auto &input = iInput.at(group_name); + + input.setShape(0, target_n.at(ifeature)); + auto tdata = input.allocate(true); + + for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) { + const auto &taginfo = (*tag_infos)[jet_n]; + const auto &features = taginfo.features(); + auto &vdata = (*tdata)[jet_n]; + if (ifeature == parT::kChargedCandidates || ifeature == parT::kChargedCandidates4Vec) + parT_tensor_filler(vdata, ifeature, features.c_pf_features, target_n_cpf); + else if (ifeature == parT::kNeutralCandidates || ifeature == parT::kNeutralCandidates4Vec) + parT_tensor_filler(vdata, ifeature, features.n_pf_features, target_n_npf); + else if (ifeature == parT::kVertices || ifeature == parT::kVertices4Vec) + parT_tensor_filler(vdata, ifeature, features.sv_features, target_n_vtx); + } + input.toServer(tdata); + } +} + +void ParticleTransformerAK4SonicJetTagsProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup, Output const &iOutput) { + edm::Handle tag_infos; + iEvent.getByToken(src_, tag_infos); + + // initialize output collection + std::vector> output_tags; + if (!tag_infos->empty()) { + auto jet_ref = tag_infos->begin()->jet(); + auto ref2prod = edm::makeRefToBaseProdFrom(jet_ref, iEvent); + for (std::size_t i = 0; i < flav_names_.size(); i++) { + output_tags.emplace_back(std::make_unique(ref2prod)); + } + } else { + for (std::size_t i = 0; i < flav_names_.size(); i++) { + output_tags.emplace_back(std::make_unique()); + } + } + if (!tag_infos->empty()) { + if (!skippedInference_) { + const auto &output1 = iOutput.begin()->second; + const auto &outputs_from_server = output1.fromServer(); + + for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) { + const auto &taginfo = (*tag_infos)[jet_n]; + const auto &jet_ref = tag_infos->at(jet_n).jet(); + + if (taginfo.features().is_filled) { + for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) { + (*(output_tags[flav_n]))[jet_ref] = outputs_from_server[jet_n][flav_n]; + } + } else { + for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) { + (*(output_tags[flav_n]))[jet_ref] = -1.0; + } + } + } + } else { + for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) { + const auto &jet_ref = tag_infos->at(jet_n).jet(); + for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) { + (*(output_tags[flav_n]))[jet_ref] = -1.0; + } + } + } + } + // put into the event + for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) { + iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]); + } +} + +//define this as a plug-in +DEFINE_FWK_MODULE(ParticleTransformerAK4SonicJetTagsProducer); diff --git a/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.cc index a6864da0c3de5..2f9548d90f880 100644 --- a/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.cc +++ b/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.cc @@ -16,6 +16,9 @@ #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h" +#include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h" +#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h" + using namespace cms::Ort; class UnifiedParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer> { @@ -42,28 +45,10 @@ class UnifiedParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDP std::vector input_names_; std::vector output_names_; - enum InputIndexes { - kChargedCandidates = 0, - kLostTracks = 1, - kNeutralCandidates = 2, - kVertices = 3, - kChargedCandidates4Vec = 4, - kLostTracks4Vec = 5, - kNeutralCandidates4Vec = 6, - kVertices4Vec = 7 - }; unsigned n_cpf_; - constexpr static unsigned n_features_cpf_ = 25; - constexpr static unsigned n_pairwise_features_cpf_ = 4; unsigned n_lt_; - constexpr static unsigned n_features_lt_ = 18; - constexpr static unsigned n_pairwise_features_lt_ = 4; unsigned n_npf_; - constexpr static unsigned n_features_npf_ = 8; - constexpr static unsigned n_pairwise_features_npf_ = 4; unsigned n_sv_; - constexpr static unsigned n_features_sv_ = 14; - constexpr static unsigned n_pairwise_features_sv_ = 4; std::vector input_sizes_; std::vector> input_shapes_; // shapes of each input group (-1 for dynamic axis) @@ -89,7 +74,7 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::Con desc.add("src", edm::InputTag("pfUnifiedParticleTransformerAK4TagInfos")); desc.add>( "input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6", "input_7", "input_8"}); - desc.add("model_path", edm::FileInPath("RecoBTag/Combined/data/UParTAK4/PUPPI/V00/UParTAK4.onnx")); + desc.add("model_path", edm::FileInPath("RecoBTag/Combined/data/UParTAK4/PUPPI/V00/modelfile/model.onnx")); desc.add>("output_names", {"softmax"}); desc.add>( "flav_names", @@ -135,14 +120,14 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::produce(edm::Event& iEven get_input_sizes(taginfo); // run prediction with dynamic batch size per event - input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)n_features_cpf_}, - {(int64_t)1, (int64_t)n_lt_, (int64_t)n_features_lt_}, - {(int64_t)1, (int64_t)n_npf_, (int64_t)n_features_npf_}, - {(int64_t)1, (int64_t)n_sv_, (int64_t)n_features_sv_}, - {(int64_t)1, (int64_t)n_cpf_, (int64_t)n_pairwise_features_cpf_}, - {(int64_t)1, (int64_t)n_lt_, (int64_t)n_pairwise_features_lt_}, - {(int64_t)1, (int64_t)n_npf_, (int64_t)n_pairwise_features_npf_}, - {(int64_t)1, (int64_t)n_sv_, (int64_t)n_pairwise_features_sv_}}; + input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)UparT::N_InputFeatures.at(UparT::kChargedCandidates)}, + {(int64_t)1, (int64_t)n_lt_, (int64_t)UparT::N_InputFeatures.at(UparT::kLostTracks)}, + {(int64_t)1, (int64_t)n_npf_, (int64_t)UparT::N_InputFeatures.at(UparT::kNeutralCandidates)}, + {(int64_t)1, (int64_t)n_sv_, (int64_t)UparT::N_InputFeatures.at(UparT::kVertices)}, + {(int64_t)1, (int64_t)n_cpf_, (int64_t)UparT::N_InputFeatures.at(UparT::kChargedCandidates4Vec)}, + {(int64_t)1, (int64_t)n_lt_, (int64_t)UparT::N_InputFeatures.at(UparT::kLostTracks4Vec)}, + {(int64_t)1, (int64_t)n_npf_, (int64_t)UparT::N_InputFeatures.at(UparT::kNeutralCandidates4Vec)}, + {(int64_t)1, (int64_t)n_sv_, (int64_t)UparT::N_InputFeatures.at(UparT::kVertices4Vec)}}; outputs = globalCache()->run(input_names_, data_, input_shapes_, output_names_, 1)[0]; assert(outputs.size() == flav_names_.size()); @@ -165,20 +150,20 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes( const auto& features = taginfo.features(); /// We require a fixed size due to an ONNX conversion issue (to be improved in the future ?) /// - n_cpf_ = (unsigned int)29; - n_lt_ = (unsigned int)5; - n_npf_ = (unsigned int)25; - n_sv_ = (unsigned int)5; + n_cpf_ = (unsigned int)UparT::n_cpf_accept; + n_lt_ = (unsigned int)UparT::n_lt_accept; + n_npf_ = (unsigned int)UparT::n_npf_accept; + n_sv_ = (unsigned int)UparT::n_sv_accept; input_sizes_ = { - n_cpf_ * n_features_cpf_, - n_lt_ * n_features_lt_, - n_npf_ * n_features_npf_, - n_sv_ * n_features_sv_, - n_cpf_ * n_pairwise_features_cpf_, - n_lt_ * n_pairwise_features_lt_, - n_npf_ * n_pairwise_features_npf_, - n_sv_ * n_pairwise_features_sv_, + n_cpf_ * UparT::N_InputFeatures.at(UparT::kChargedCandidates), + n_lt_ * UparT::N_InputFeatures.at(UparT::kLostTracks), + n_npf_ * UparT::N_InputFeatures.at(UparT::kNeutralCandidates), + n_sv_ * UparT::N_InputFeatures.at(UparT::kVertices), + n_cpf_ * UparT::N_InputFeatures.at(UparT::kChargedCandidates4Vec), + n_lt_ * UparT::N_InputFeatures.at(UparT::kLostTracks4Vec), + n_npf_ * UparT::N_InputFeatures.at(UparT::kNeutralCandidates4Vec), + n_sv_ * UparT::N_InputFeatures.at(UparT::kVertices4Vec), }; // init data storage data_.clear(); @@ -191,167 +176,30 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes( void UnifiedParticleTransformerAK4ONNXJetTagsProducer::make_inputs( btagbtvdeep::UnifiedParticleTransformerAK4Features features) { - float* ptr = nullptr; const float* start = nullptr; unsigned offset = 0; - // c_pf candidates auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_); - for (std::size_t c_pf_n = 0; c_pf_n < max_c_pf_n; c_pf_n++) { - const auto& c_pf_features = features.c_pf_features.at(c_pf_n); - ptr = &data_[kChargedCandidates][offset + c_pf_n * n_features_cpf_]; - start = ptr; - *ptr = c_pf_features.btagPf_trackEtaRel; - *(++ptr) = c_pf_features.btagPf_trackPtRel; - *(++ptr) = c_pf_features.btagPf_trackPPar; - *(++ptr) = c_pf_features.btagPf_trackDeltaR; - *(++ptr) = c_pf_features.btagPf_trackPParRatio; - *(++ptr) = c_pf_features.btagPf_trackSip2dVal; - *(++ptr) = c_pf_features.btagPf_trackSip2dSig; - *(++ptr) = c_pf_features.btagPf_trackSip3dVal; - *(++ptr) = c_pf_features.btagPf_trackSip3dSig; - *(++ptr) = c_pf_features.btagPf_trackJetDistVal; - *(++ptr) = c_pf_features.ptrel; - *(++ptr) = c_pf_features.drminsv; - *(++ptr) = c_pf_features.vtx_ass; - *(++ptr) = c_pf_features.puppiw; - *(++ptr) = c_pf_features.chi2; - *(++ptr) = c_pf_features.quality; - *(++ptr) = c_pf_features.charge; - *(++ptr) = c_pf_features.dz; - *(++ptr) = c_pf_features.btagPf_trackDecayLen; - *(++ptr) = c_pf_features.HadFrac; - *(++ptr) = c_pf_features.CaloFrac; - *(++ptr) = c_pf_features.pdgID; - *(++ptr) = c_pf_features.lostInnerHits; - *(++ptr) = c_pf_features.numberOfPixelHits; - *(++ptr) = c_pf_features.numberOfStripHits; - - assert(start + n_features_cpf_ - 1 == ptr); - } - - // n_lt candidates auto max_lt_n = std::min(features.lt_features.size(), (std::size_t)n_lt_); - for (std::size_t lt_n = 0; lt_n < max_lt_n; lt_n++) { - const auto& lt_features = features.lt_features.at(lt_n); - ptr = &data_[kLostTracks][offset + lt_n * n_features_lt_]; - start = ptr; - *ptr = lt_features.btagPf_trackEtaRel; - *(++ptr) = lt_features.btagPf_trackPtRel; - *(++ptr) = lt_features.btagPf_trackPPar; - *(++ptr) = lt_features.btagPf_trackDeltaR; - *(++ptr) = lt_features.btagPf_trackPParRatio; - *(++ptr) = lt_features.btagPf_trackSip2dVal; - *(++ptr) = lt_features.btagPf_trackSip2dSig; - *(++ptr) = lt_features.btagPf_trackSip3dVal; - *(++ptr) = lt_features.btagPf_trackSip3dSig; - *(++ptr) = lt_features.btagPf_trackJetDistVal; - *(++ptr) = lt_features.drminsv; - *(++ptr) = lt_features.charge; - *(++ptr) = lt_features.puppiw; - *(++ptr) = lt_features.chi2; - *(++ptr) = lt_features.quality; - *(++ptr) = lt_features.lostInnerHits; - *(++ptr) = lt_features.numberOfPixelHits; - *(++ptr) = lt_features.numberOfStripHits; - assert(start + n_features_lt_ - 1 == ptr); - } - - // n_pf candidates auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_); - for (std::size_t n_pf_n = 0; n_pf_n < max_n_pf_n; n_pf_n++) { - const auto& n_pf_features = features.n_pf_features.at(n_pf_n); - ptr = &data_[kNeutralCandidates][offset + n_pf_n * n_features_npf_]; - start = ptr; - *ptr = n_pf_features.ptrel; - *(++ptr) = n_pf_features.etarel; - *(++ptr) = n_pf_features.phirel; - *(++ptr) = n_pf_features.deltaR; - *(++ptr) = n_pf_features.isGamma; - *(++ptr) = n_pf_features.hadFrac; - *(++ptr) = n_pf_features.drminsv; - *(++ptr) = n_pf_features.puppiw; - assert(start + n_features_npf_ - 1 == ptr); - } - - // sv candidates auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_); - for (std::size_t sv_n = 0; sv_n < max_sv_n; sv_n++) { - const auto& sv_features = features.sv_features.at(sv_n); - ptr = &data_[kVertices][offset + sv_n * n_features_sv_]; - start = ptr; - *ptr = sv_features.pt; - *(++ptr) = sv_features.deltaR; - *(++ptr) = sv_features.mass; - *(++ptr) = sv_features.etarel; - *(++ptr) = sv_features.phirel; - *(++ptr) = sv_features.ntracks; - *(++ptr) = sv_features.chi2; - *(++ptr) = sv_features.normchi2; - *(++ptr) = sv_features.dxy; - *(++ptr) = sv_features.dxysig; - *(++ptr) = sv_features.d3d; - *(++ptr) = sv_features.d3dsig; - *(++ptr) = sv_features.costhetasvpv; - *(++ptr) = sv_features.enratio; - assert(start + n_features_sv_ - 1 == ptr); - } + // c_pf candidates + UparT_tensor_filler(data_, UparT::kChargedCandidates, features.c_pf_features, max_c_pf_n, start, offset); + // lt candidates + UparT_tensor_filler(data_, UparT::kLostTracks, features.lt_features, max_lt_n, start, offset); + // n_pf candidates + UparT_tensor_filler(data_, UparT::kNeutralCandidates, features.n_pf_features, max_n_pf_n, start, offset); + // sv candidates + UparT_tensor_filler(data_, UparT::kVertices, features.sv_features, max_sv_n, start, offset); // cpf pairwise features (4-vectors) - auto max_cpf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_); - for (std::size_t cpf_n = 0; cpf_n < max_cpf_n; cpf_n++) { - const auto& cpf_pairwise_features = features.c_pf_features.at(cpf_n); - ptr = &data_[kChargedCandidates4Vec][offset + cpf_n * n_pairwise_features_cpf_]; - start = ptr; - *ptr = cpf_pairwise_features.px; - *(++ptr) = cpf_pairwise_features.py; - *(++ptr) = cpf_pairwise_features.pz; - *(++ptr) = cpf_pairwise_features.e; - - assert(start + n_pairwise_features_cpf_ - 1 == ptr); - } - + UparT_tensor_filler(data_, UparT::kChargedCandidates4Vec, features.c_pf_features, max_c_pf_n, start, offset); // lt pairwise features (4-vectors) specific case requiring (pt,eta,phi,e) - auto max_lt_N = std::min(features.lt_features.size(), (std::size_t)n_lt_); - for (std::size_t lt_N = 0; lt_N < max_lt_N; lt_N++) { - const auto& lt_pairwise_features = features.lt_features.at(lt_N); - ptr = &data_[kLostTracks4Vec][offset + lt_N * n_pairwise_features_lt_]; - start = ptr; - *ptr = lt_pairwise_features.pt; - *(++ptr) = lt_pairwise_features.eta; - *(++ptr) = lt_pairwise_features.phi; - *(++ptr) = lt_pairwise_features.e; - - assert(start + n_pairwise_features_lt_ - 1 == ptr); - } - + UparT_tensor_filler(data_, UparT::kLostTracks4Vec, features.lt_features, max_lt_n, start, offset); // npf pairwise features (4-vectors) - auto max_npf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_); - for (std::size_t npf_n = 0; npf_n < max_npf_n; npf_n++) { - const auto& npf_pairwise_features = features.n_pf_features.at(npf_n); - ptr = &data_[kNeutralCandidates4Vec][offset + npf_n * n_pairwise_features_npf_]; - start = ptr; - *ptr = npf_pairwise_features.px; - *(++ptr) = npf_pairwise_features.py; - *(++ptr) = npf_pairwise_features.pz; - *(++ptr) = npf_pairwise_features.e; - - assert(start + n_pairwise_features_npf_ - 1 == ptr); - } - + UparT_tensor_filler(data_, UparT::kNeutralCandidates4Vec, features.n_pf_features, max_n_pf_n, start, offset); // sv pairwise features (4-vectors) - auto max_sv_N = std::min(features.sv_features.size(), (std::size_t)n_sv_); - for (std::size_t sv_N = 0; sv_N < max_sv_N; sv_N++) { - const auto& sv_pairwise_features = features.sv_features.at(sv_N); - ptr = &data_[kVertices4Vec][offset + sv_N * n_pairwise_features_sv_]; - start = ptr; - *ptr = sv_pairwise_features.px; - *(++ptr) = sv_pairwise_features.py; - *(++ptr) = sv_pairwise_features.pz; - *(++ptr) = sv_pairwise_features.e; - - assert(start + n_pairwise_features_sv_ - 1 == ptr); - } + UparT_tensor_filler(data_, UparT::kVertices4Vec, features.sv_features, max_sv_n, start, offset); } //define this as a plug-in diff --git a/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc new file mode 100644 index 0000000000000..bc7810f159dd4 --- /dev/null +++ b/RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4SonicJetTagsProducer.cc @@ -0,0 +1,202 @@ +#include "FWCore/Framework/interface/Frameworkfwd.h" + +#include "FWCore/Framework/interface/Event.h" +#include "FWCore/Framework/interface/MakerMacros.h" + +#include "FWCore/Framework/interface/makeRefToBaseProdFrom.h" + +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/Utilities/interface/StreamID.h" + +#include "DataFormats/BTauReco/interface/JetTag.h" + +#include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4TagInfo.h" +#include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4Features.h" + +#include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonData.h" + +#include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h" +#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h" + +class UnifiedParticleTransformerAK4SonicJetTagsProducer : public TritonEDProducer<> { +public: + explicit UnifiedParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet&); + ~UnifiedParticleTransformerAK4SonicJetTagsProducer() override; + + void acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) override; + + void produce(edm::Event &iEvent, edm::EventSetup const &iSetup, Output const &iOutput) override; + static void fillDescriptions(edm::ConfigurationDescriptions&); + +private: + typedef std::vector TagInfoCollection; + typedef reco::JetTagCollection JetTagCollection; + + const edm::EDGetTokenT src_; + std::vector flav_names_; + std::vector input_names_; + std::vector output_names_; + + bool skippedInference_ = false; +}; + +UnifiedParticleTransformerAK4SonicJetTagsProducer::UnifiedParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet& iConfig) + : TritonEDProducer<>(iConfig), + src_(consumes(iConfig.getParameter("src"))), + flav_names_(iConfig.getParameter>("flav_names")), + input_names_(iConfig.getParameter>("input_names")), + output_names_(iConfig.getParameter>("output_names")) { + // get output names from flav_names + for (const auto& flav_name : flav_names_) { + produces(flav_name); + } +} + +UnifiedParticleTransformerAK4SonicJetTagsProducer::~UnifiedParticleTransformerAK4SonicJetTagsProducer() {} + +void UnifiedParticleTransformerAK4SonicJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) { + // pfUnifiedParticleTransformerAK4JetTags + edm::ParameterSetDescription desc; + TritonClient::fillPSetDescription(desc); + desc.add("src", edm::InputTag("pfUnifiedParticleTransformerAK4TagInfos")); + desc.add>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6", "input_7", "input_8"}); + desc.add>("output_names", {"softmax"}); + desc.add>( + "flav_names", + std::vector{"probb", "probbb", "problepb", "probc", "probs", + "probu", "probd", "probg", "probele", "probmu", + "probtaup1h0p", "probtaup1h1p", "probtaup1h2p", "probtaup3h0p", "probtaup3h1p", + "probtaum1h0p", "probtaum1h1p", "probtaum1h2p", "probtaum3h0p", "probtaum3h1p", + "ptcorr", "ptreshigh", "ptreslow", "ptnu", "probemudata", + "probemumc", "probdimudata", "probdimumc", "probmutaudata", "probmutaumc"}); + + descriptions.add("pfUnifiedParticleTransformerAK4SonicJetTags", desc); +} + +void UnifiedParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) { + edm::Handle tag_infos; + iEvent.getByToken(src_, tag_infos); + client_->setBatchSize(tag_infos->size()); + skippedInference_ = false; + if (tag_infos->empty()) return; + + // Find the max n_cpf, n_npf and n_vtx among all the jets in an event. + + unsigned int max_n_cpf_counter = 0; + unsigned int max_n_lt_counter = 0; + unsigned int max_n_npf_counter = 0; + unsigned int max_n_vtx_counter = 0; + for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) { + max_n_cpf_counter = std::max(max_n_cpf_counter, + static_cast(((*tag_infos)[jet_n]).features().c_pf_features.size())); + max_n_lt_counter = std::max(max_n_lt_counter, + static_cast(((*tag_infos)[jet_n]).features().lt_features.size())); + max_n_npf_counter = std::max(max_n_npf_counter, + static_cast(((*tag_infos)[jet_n]).features().n_pf_features.size())); + max_n_vtx_counter = std::max(max_n_vtx_counter, + static_cast(((*tag_infos)[jet_n]).features().sv_features.size())); + } + + // If an event has no jet, or all jets has zero n_cpf, n_lt, n_npf and n_vtx, the inference is skipped. + if (max_n_cpf_counter == 0 && max_n_lt_counter == 0 && max_n_npf_counter == 0 && max_n_vtx_counter == 0) { + client_->setBatchSize(0); + skippedInference_ = true; + return; + } + + // all the jets in the same event will fill up the same amount of n_cpf, n_npf, n_vtx and send to server + const unsigned int target_n_cpf = (unsigned int)UparT::n_cpf_accept; + const unsigned int target_n_lt = (unsigned int)UparT::n_lt_accept; + const unsigned int target_n_npf = (unsigned int)UparT::n_npf_accept; + const unsigned int target_n_vtx = (unsigned int)UparT::n_sv_accept; + + const std::map target_n { + {UparT::kChargedCandidates, target_n_cpf}, + {UparT::kLostTracks, target_n_lt}, + {UparT::kNeutralCandidates, target_n_npf}, + {UparT::kVertices, target_n_vtx}, + {UparT::kChargedCandidates4Vec, target_n_cpf}, + {UparT::kLostTracks4Vec, target_n_lt}, + {UparT::kNeutralCandidates4Vec, target_n_npf}, + {UparT::kVertices4Vec, target_n_vtx} + }; + + // loop through all groups of features + for (UparT::InputFeatures ifeature = UparT::kBegin; ifeature != UparT::kEnd; ifeature = static_cast(ifeature+1)) { + const auto &group_name = input_names_[ifeature]; + auto &input = iInput.at(group_name); + + input.setShape(0, target_n.at(ifeature)); + auto tdata = input.allocate(true); + + for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) { + const auto &taginfo = (*tag_infos)[jet_n]; + const auto &features = taginfo.features(); + auto &vdata = (*tdata)[jet_n]; + if (ifeature == UparT::kChargedCandidates || ifeature == UparT::kChargedCandidates4Vec) + UparT_tensor_filler(vdata, ifeature, features.c_pf_features, target_n_cpf); + else if (ifeature == UparT::kLostTracks || ifeature == UparT::kLostTracks4Vec) + UparT_tensor_filler(vdata, ifeature, features.lt_features, target_n_lt); + else if (ifeature == UparT::kNeutralCandidates || ifeature == UparT::kNeutralCandidates4Vec) + UparT_tensor_filler(vdata, ifeature, features.n_pf_features, target_n_npf); + else if (ifeature == UparT::kVertices || ifeature == UparT::kVertices4Vec) + UparT_tensor_filler(vdata, ifeature, features.sv_features, target_n_vtx); + } + input.toServer(tdata); + } +} + +void UnifiedParticleTransformerAK4SonicJetTagsProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup, Output const &iOutput) { + edm::Handle tag_infos; + iEvent.getByToken(src_, tag_infos); + + // initialize output collection + std::vector> output_tags; + if (!tag_infos->empty()) { + auto jet_ref = tag_infos->begin()->jet(); + auto ref2prod = edm::makeRefToBaseProdFrom(jet_ref, iEvent); + for (std::size_t i = 0; i < flav_names_.size(); i++) { + output_tags.emplace_back(std::make_unique(ref2prod)); + } + } else { + for (std::size_t i = 0; i < flav_names_.size(); i++) { + output_tags.emplace_back(std::make_unique()); + } + } + if (!tag_infos->empty()) { + if (!skippedInference_) { + const auto &output1 = iOutput.begin()->second; + const auto &outputs_from_server = output1.fromServer(); + + for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) { + const auto &taginfo = (*tag_infos)[jet_n]; + const auto &jet_ref = tag_infos->at(jet_n).jet(); + + if (taginfo.features().is_filled) { + for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) { + (*(output_tags[flav_n]))[jet_ref] = outputs_from_server[jet_n][flav_n]; + } + } else { + for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) { + (*(output_tags[flav_n]))[jet_ref] = -1.0; + } + } + } + } else { + for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) { + const auto &jet_ref = tag_infos->at(jet_n).jet(); + for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) { + (*(output_tags[flav_n]))[jet_ref] = -1.0; + } + } + } + } + // put into the event + for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) { + iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]); + } +} + +//define this as a plug-in +DEFINE_FWK_MODULE(UnifiedParticleTransformerAK4SonicJetTagsProducer); diff --git a/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py b/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py index 9df64c14f3406..b721636164802 100644 --- a/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py +++ b/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py @@ -6,6 +6,23 @@ from RecoBTag.ONNXRuntime.pfParticleTransformerAK4DiscriminatorsJetTags_cfi import pfParticleTransformerAK4DiscriminatorsJetTags from CommonTools.PileupAlgos.Puppi_cff import puppi from CommonTools.RecoAlgos.primaryVertexAssociation_cfi import primaryVertexAssociation +from RecoBTag.ONNXRuntime.pfParticleTransformerAK4SonicJetTags_cfi import pfParticleTransformerAK4SonicJetTags as _pfParticleTransformerAK4SonicJetTags +from Configuration.ProcessModifiers.particleTransformerAK4SonicTriton_cff import particleTransformerAK4SonicTriton + +particleTransformerAK4SonicTriton.toReplaceWith(pfParticleTransformerAK4JetTags, _pfParticleTransformerAK4SonicJetTags.clone( + Client = cms.PSet( + timeout = cms.untracked.uint32(300), + mode = cms.string("Async"), + modelName = cms.string("particletransformer_AK4"), # double check + modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particletransformer_AK4/config.pbtxt"), # this is SONIC, not currently in the CMSSW, so the models/ will be copied to this location privately + modelVersion = cms.string(""), + verbose = cms.untracked.bool(False), + allowedTries = cms.untracked.uint32(0), + useSharedMemory = cms.untracked.bool(True), + compression = cms.untracked.string(""), + ), + flav_names = pfParticleTransformerAK4JetTags.flav_names, +)) # declare all the discriminators # probs diff --git a/RecoBTag/ONNXRuntime/python/pfUnifiedParticleTransformerAK4_cff.py b/RecoBTag/ONNXRuntime/python/pfUnifiedParticleTransformerAK4_cff.py index aee089f50ef68..b9584d00401bc 100644 --- a/RecoBTag/ONNXRuntime/python/pfUnifiedParticleTransformerAK4_cff.py +++ b/RecoBTag/ONNXRuntime/python/pfUnifiedParticleTransformerAK4_cff.py @@ -6,6 +6,24 @@ from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4DiscriminatorsJetTags_cfi import pfUnifiedParticleTransformerAK4DiscriminatorsJetTags from CommonTools.PileupAlgos.Puppi_cff import puppi from CommonTools.RecoAlgos.primaryVertexAssociation_cfi import primaryVertexAssociation +from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4SonicJetTags_cfi import pfUnifiedParticleTransformerAK4SonicJetTags as _pfUnifiedParticleTransformerAK4SonicJetTags +from Configuration.ProcessModifiers.unifiedparticleTransformerAK4SonicTriton_cff import unifiedparticleTransformerAK4SonicTriton + +unifiedparticleTransformerAK4SonicTriton.toReplaceWith(pfUnifiedParticleTransformerAK4JetTags, _pfUnifiedParticleTransformerAK4SonicJetTags.clone( + Client = cms.PSet( + timeout = cms.untracked.uint32(300), + mode = cms.string("Async"), + modelName = cms.string("unifiedparticletransformer_AK4"), # double check + modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/unifiedparticletransformer_AK4/config.pbtxt"), # this is SONIC, not currently in the CMSSW, so the models/ will be copied to this location privately + modelVersion = cms.string(""), + verbose = cms.untracked.bool(False), + allowedTries = cms.untracked.uint32(0), + useSharedMemory = cms.untracked.bool(True), + compression = cms.untracked.string(""), + ), + flav_names = pfUnifiedParticleTransformerAK4JetTags.flav_names, +)) + # declare all the discriminators # probs diff --git a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc index da49b1b0119d7..3a5f0357ea6fc 100644 --- a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc +++ b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc @@ -138,4 +138,197 @@ namespace btagbtvdeep { *(++ptr) = neighbourTrack_features.dphi_PCAjetDirs; } + std::vector inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputFeatures ifeature) { + std::vector inputs; + if (ifeature == parT::kChargedCandidates) { + inputs.push_back(c_pf_features.btagPf_trackEtaRel); + inputs.push_back(c_pf_features.btagPf_trackPtRel); + inputs.push_back(c_pf_features.btagPf_trackPPar); + inputs.push_back(c_pf_features.btagPf_trackDeltaR); + inputs.push_back(c_pf_features.btagPf_trackPParRatio); + inputs.push_back(c_pf_features.btagPf_trackSip2dVal); + inputs.push_back(c_pf_features.btagPf_trackSip2dSig); + inputs.push_back(c_pf_features.btagPf_trackSip3dVal); + inputs.push_back(c_pf_features.btagPf_trackSip3dSig); + inputs.push_back(c_pf_features.btagPf_trackJetDistVal); + inputs.push_back(c_pf_features.ptrel); + inputs.push_back(c_pf_features.drminsv); + inputs.push_back(c_pf_features.vtx_ass); + inputs.push_back(c_pf_features.puppiw); + inputs.push_back(c_pf_features.chi2); + inputs.push_back(c_pf_features.quality); + } + else if (ifeature == parT::kChargedCandidates4Vec) { + inputs.push_back(c_pf_features.px); + inputs.push_back(c_pf_features.py); + inputs.push_back(c_pf_features.pz); + inputs.push_back(c_pf_features.e); + } + return inputs; + } + + std::vector inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputFeatures ifeature) { + std::vector inputs; + if (ifeature == parT::kNeutralCandidates) { + inputs.push_back(n_pf_features.ptrel); + inputs.push_back(n_pf_features.etarel); + inputs.push_back(n_pf_features.phirel); + inputs.push_back(n_pf_features.deltaR); + inputs.push_back(n_pf_features.isGamma); + inputs.push_back(n_pf_features.hadFrac); + inputs.push_back(n_pf_features.drminsv); + inputs.push_back(n_pf_features.puppiw); + } + else if (ifeature == parT::kNeutralCandidates4Vec) { + inputs.push_back(n_pf_features.px); + inputs.push_back(n_pf_features.py); + inputs.push_back(n_pf_features.pz); + inputs.push_back(n_pf_features.e); + } + return inputs; + } + + std::vector inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature) { + std::vector inputs; + if (ifeature == parT::kVertices) { + inputs.push_back(sv_features.pt); + inputs.push_back(sv_features.deltaR); + inputs.push_back(sv_features.mass); + inputs.push_back(sv_features.etarel); + inputs.push_back(sv_features.phirel); + inputs.push_back(sv_features.ntracks); + inputs.push_back(sv_features.chi2); + inputs.push_back(sv_features.normchi2); + inputs.push_back(sv_features.dxy); + inputs.push_back(sv_features.dxysig); + inputs.push_back(sv_features.d3d); + inputs.push_back(sv_features.d3dsig); + inputs.push_back(sv_features.costhetasvpv); + inputs.push_back(sv_features.enratio); + } + else if (ifeature == parT::kVertices4Vec) { + inputs.push_back(sv_features.px); + inputs.push_back(sv_features.py); + inputs.push_back(sv_features.pz); + inputs.push_back(sv_features.e); + } + return inputs; + } + + std::vector inputs_UparT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, UparT::InputFeatures ifeature) { + std::vector inputs; + if (ifeature == UparT::kChargedCandidates) { + inputs.push_back(c_pf_features.btagPf_trackEtaRel); + inputs.push_back(c_pf_features.btagPf_trackPtRel); + inputs.push_back(c_pf_features.btagPf_trackPPar); + inputs.push_back(c_pf_features.btagPf_trackDeltaR); + inputs.push_back(c_pf_features.btagPf_trackPParRatio); + inputs.push_back(c_pf_features.btagPf_trackSip2dVal); + inputs.push_back(c_pf_features.btagPf_trackSip2dSig); + inputs.push_back(c_pf_features.btagPf_trackSip3dVal); + inputs.push_back(c_pf_features.btagPf_trackSip3dSig); + inputs.push_back(c_pf_features.btagPf_trackJetDistVal); + inputs.push_back(c_pf_features.ptrel); + inputs.push_back(c_pf_features.drminsv); + inputs.push_back(c_pf_features.vtx_ass); + inputs.push_back(c_pf_features.puppiw); + inputs.push_back(c_pf_features.chi2); + inputs.push_back(c_pf_features.quality); + inputs.push_back(c_pf_features.charge); + inputs.push_back(c_pf_features.dz); + inputs.push_back(c_pf_features.btagPf_trackDecayLen); + inputs.push_back(c_pf_features.HadFrac); + inputs.push_back(c_pf_features.CaloFrac); + inputs.push_back(c_pf_features.pdgID); + inputs.push_back(c_pf_features.lostInnerHits); + inputs.push_back(c_pf_features.numberOfPixelHits); + inputs.push_back(c_pf_features.numberOfStripHits); + } + else if (ifeature == UparT::kChargedCandidates4Vec) { + inputs.push_back(c_pf_features.px); + inputs.push_back(c_pf_features.py); + inputs.push_back(c_pf_features.pz); + inputs.push_back(c_pf_features.e); + } + return inputs; + } + + std::vector inputs_UparT(const btagbtvdeep::LostTracksFeatures& lt_features, UparT::InputFeatures ifeature) { + std::vector inputs; + if (ifeature == UparT::kLostTracks) { + inputs.push_back(lt_features.btagPf_trackEtaRel); + inputs.push_back(lt_features.btagPf_trackPtRel); + inputs.push_back(lt_features.btagPf_trackPPar); + inputs.push_back(lt_features.btagPf_trackDeltaR); + inputs.push_back(lt_features.btagPf_trackPParRatio); + inputs.push_back(lt_features.btagPf_trackSip2dVal); + inputs.push_back(lt_features.btagPf_trackSip2dSig); + inputs.push_back(lt_features.btagPf_trackSip3dVal); + inputs.push_back(lt_features.btagPf_trackSip3dSig); + inputs.push_back(lt_features.btagPf_trackJetDistVal); + inputs.push_back(lt_features.drminsv); + inputs.push_back(lt_features.charge); + inputs.push_back(lt_features.puppiw); + inputs.push_back(lt_features.chi2); + inputs.push_back(lt_features.quality); + inputs.push_back(lt_features.lostInnerHits); + inputs.push_back(lt_features.numberOfPixelHits); + inputs.push_back(lt_features.numberOfStripHits); + } + else if (ifeature == UparT::kLostTracks4Vec) { + inputs.push_back(lt_features.pt); + inputs.push_back(lt_features.eta); + inputs.push_back(lt_features.phi); + inputs.push_back(lt_features.e); + } + return inputs; + } + + std::vector inputs_UparT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, UparT::InputFeatures ifeature) { + std::vector inputs; + if (ifeature == UparT::kNeutralCandidates) { + inputs.push_back(n_pf_features.ptrel); + inputs.push_back(n_pf_features.etarel); + inputs.push_back(n_pf_features.phirel); + inputs.push_back(n_pf_features.deltaR); + inputs.push_back(n_pf_features.isGamma); + inputs.push_back(n_pf_features.hadFrac); + inputs.push_back(n_pf_features.drminsv); + inputs.push_back(n_pf_features.puppiw); + } + else if (ifeature == UparT::kNeutralCandidates4Vec) { + inputs.push_back(n_pf_features.px); + inputs.push_back(n_pf_features.py); + inputs.push_back(n_pf_features.pz); + inputs.push_back(n_pf_features.e); + } + return inputs; + } + + std::vector inputs_UparT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, UparT::InputFeatures ifeature) { + std::vector inputs; + if (ifeature == UparT::kVertices) { + inputs.push_back(sv_features.pt); + inputs.push_back(sv_features.deltaR); + inputs.push_back(sv_features.mass); + inputs.push_back(sv_features.etarel); + inputs.push_back(sv_features.phirel); + inputs.push_back(sv_features.ntracks); + inputs.push_back(sv_features.chi2); + inputs.push_back(sv_features.normchi2); + inputs.push_back(sv_features.dxy); + inputs.push_back(sv_features.dxysig); + inputs.push_back(sv_features.d3d); + inputs.push_back(sv_features.d3dsig); + inputs.push_back(sv_features.costhetasvpv); + inputs.push_back(sv_features.enratio); + } + else if (ifeature == UparT::kVertices4Vec) { + inputs.push_back(sv_features.px); + inputs.push_back(sv_features.py); + inputs.push_back(sv_features.pz); + inputs.push_back(sv_features.e); + } + return inputs; + } } // namespace btagbtvdeep