diff --git a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py index 396fcefd84a79..f349808c487f7 100644 --- a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py +++ b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py @@ -5,6 +5,8 @@ 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 # collect all SonicTriton-related process modifiers here -allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,deepTauSonicTriton) +allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,particleNetPTSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton) + 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/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..98faeda8403db 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_configs.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_configs.h @@ -1,6 +1,7 @@ #ifndef RecoBTag_ONNXRuntime_tensor_configs_h #define RecoBTag_ONNXRuntime_tensor_configs_h +#include namespace deepflavour { constexpr unsigned n_features_global = 15; @@ -28,4 +29,41 @@ namespace deepvertex { } // namespace deepvertex +namespace parT { + + enum InputFeatures { + kBegin = 0, + kChargedCandidates = kBegin, + kNeutralCandidates = 1, + kVertices = 2, + kChargedCandidates4Vec = 3, + kNeutralCandidates4Vec = 4, + kVertices4Vec = 5, + kEnd = 6 + }; + + inline constexpr unsigned n_cpf_accept = 25; + inline constexpr unsigned n_npf_accept = 25; + inline constexpr unsigned n_sv_accept = 5; + + constexpr std::array N_InputFeatures{{ + 16, // kChargedCandidates + 8, // kNeutralCandidates + 14, // kVertices + 4, // kChargedCandidates4Vec + 4, // kNeutralCandidates4Vec + 4, // kVertices4Vec + }}; + + constexpr std::array N_AcceptedFeatures{{ + n_cpf_accept, // kChargedCandidates + n_npf_accept, // kNeutralCandidates + n_sv_accept, // kVertices + n_cpf_accept, // kChargedCandidates4Vec + n_npf_accept, // kNeutralCandidates4Vec + n_sv_accept, // kVertices4Vec + }}; + +} // namespace parT + #endif diff --git a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h index bb1a25a2d0f0b..db3facc2338d1 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h @@ -2,6 +2,9 @@ #define RecoBTag_ONNXRuntime_tensor_fillers_h #include "DataFormats/BTauReco/interface/DeepFlavourTagInfo.h" +#include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h" +#include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h" +#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h" namespace btagbtvdeep { @@ -19,6 +22,53 @@ 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); + + 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 + } + } // namespace btagbtvdeep #endif diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc index b27a9b397e520..ea5fc6b0020f4 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> { @@ -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::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::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::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,26 @@ void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes( } void ParticleTransformerAK4ONNXJetTagsProducer::make_inputs(btagbtvdeep::ParticleTransformerAK4Features features) { - float* ptr = nullptr; + //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..2a33e976e61f0 --- /dev/null +++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc @@ -0,0 +1,191 @@ +#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/python/pfParticleTransformerAK4_cff.py b/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py index 0fb1c8ceac050..4a7fd51782280 100644 --- a/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py +++ b/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py @@ -4,8 +4,26 @@ from RecoBTag.ONNXRuntime.pfParticleTransformerAK4JetTags_cfi import pfParticleTransformerAK4JetTags from RecoBTag.ONNXRuntime.pfParticleTransformerAK4DiscriminatorsJetTags_cfi import pfParticleTransformerAK4DiscriminatorsJetTags +from RecoBTag.ONNXRuntime.pfParticleTransformerAK4SonicJetTags_cfi import pfParticleTransformerAK4SonicJetTags as _pfParticleTransformerAK4SonicJetTags from CommonTools.PileupAlgos.Puppi_cff import puppi from CommonTools.RecoAlgos.primaryVertexAssociation_cfi import primaryVertexAssociation +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"), + modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particletransformer_AK4/config.pbtxt"), + 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/src/tensor_fillers.cc b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc index da49b1b0119d7..9ecdf0302b45a 100644 --- a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc +++ b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc @@ -138,4 +138,81 @@ 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; + } + } // namespace btagbtvdeep