From fcb4d88da4b0922351a8148b9756d30d4ddd2371 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Tue, 2 Apr 2024 14:31:50 -0400 Subject: [PATCH 1/9] add SONIC-ize particleTransformer --- .../python/allSonicTriton_cff.py | 3 +- .../particleTransformerAK4SonicTriton_cff.py | 3 + ...ticleTransformerAK4SonicJetTagsProducer.cc | 306 ++++++++++++++++++ .../python/pfParticleTransformerAK4_cff.py | 18 ++ 4 files changed, 329 insertions(+), 1 deletion(-) create mode 100644 Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py create mode 100644 RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc diff --git a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py index 396fcefd84a79..27291329e6066 100644 --- a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py +++ b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py @@ -7,4 +7,5 @@ from Configuration.ProcessModifiers.deepTauSonicTriton_cff import deepTauSonicTriton # 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/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc new file mode 100644 index 0000000000000..989ebb95a5071 --- /dev/null +++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc @@ -0,0 +1,306 @@ +#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" + +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_; + + enum InputIndexes { + kChargedCandidates = 0, + kNeutralCandidates = 1, + kVertices = 2, + kChargedCandidates4Vec = 3, + kNeutralCandidates4Vec = 4, + kVertices4Vec = 5 + }; + + constexpr static unsigned n_features_cpf_ = 16; + constexpr static unsigned n_pairwise_features_cpf_ = 4; + constexpr static unsigned n_features_npf_ = 8; + constexpr static unsigned n_pairwise_features_npf_ = 4; + constexpr static unsigned n_features_sv_ = 14; + constexpr static unsigned n_pairwise_features_sv_ = 4; + bool skippedInference_ = false; + bool padding_ = true; +}; + +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; + unsigned int max_n_cpf = 0; + unsigned int max_n_npf = 0; + unsigned int max_n_vtx = 0; + + // Find the max n_cpf, n_npf and n_vtx among all the jets in an event. + for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) { + max_n_cpf = std::max(max_n_cpf, + static_cast(((*tag_infos)[jet_n]).features().c_pf_features.size())); + max_n_npf = std::max(max_n_npf, + static_cast(((*tag_infos)[jet_n]).features().n_pf_features.size())); + max_n_vtx = std::max(max_n_vtx, + 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 == 0 && max_n_npf == 0 && max_n_vtx == 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 + max_n_cpf = std::clamp(max_n_cpf, (unsigned int)0, (unsigned int)25); + max_n_npf = std::clamp(max_n_npf, (unsigned int)0, (unsigned int)25); + max_n_vtx = std::clamp(max_n_vtx, (unsigned int)0, (unsigned int)5); + + for (unsigned igroup = 0; igroup < input_names_.size(); ++igroup) { + const auto &group_name = input_names_[igroup]; + auto &input = iInput.at(group_name); + unsigned target = 0; + + if (igroup == kChargedCandidates || igroup == kChargedCandidates4Vec) target = std::max((unsigned int)1, max_n_cpf); + else if (igroup == kNeutralCandidates || igroup == kNeutralCandidates4Vec) target = std::max((unsigned int)1, max_n_npf); + else if (igroup == kVertices || igroup == kVertices4Vec) target = std::max((unsigned int)1, max_n_vtx); + + input.setShape(0, target); + 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]; + // Loop through the n cpf, and in the case that n_cpf is smaller than max_n_cpf, add padding values to all features + if (igroup == kChargedCandidates) { + unsigned int n_cpf = features.c_pf_features.size(); + n_cpf = std::clamp(n_cpf, (unsigned int)0, (unsigned int)25); + for (unsigned int count = 0; count < n_cpf; count++) { + vdata.push_back(features.c_pf_features.at(count).btagPf_trackEtaRel); + vdata.push_back(features.c_pf_features.at(count).btagPf_trackPtRel); + vdata.push_back(features.c_pf_features.at(count).btagPf_trackPPar); + vdata.push_back(features.c_pf_features.at(count).btagPf_trackDeltaR); + vdata.push_back(features.c_pf_features.at(count).btagPf_trackPParRatio); + vdata.push_back(features.c_pf_features.at(count).btagPf_trackSip2dVal); + vdata.push_back(features.c_pf_features.at(count).btagPf_trackSip2dSig); + vdata.push_back(features.c_pf_features.at(count).btagPf_trackSip3dVal); + vdata.push_back(features.c_pf_features.at(count).btagPf_trackSip3dSig); + vdata.push_back(features.c_pf_features.at(count).btagPf_trackJetDistVal); + vdata.push_back(features.c_pf_features.at(count).ptrel); + vdata.push_back(features.c_pf_features.at(count).drminsv); + vdata.push_back(features.c_pf_features.at(count).vtx_ass); + vdata.push_back(features.c_pf_features.at(count).puppiw); + vdata.push_back(features.c_pf_features.at(count).chi2); + vdata.push_back(features.c_pf_features.at(count).quality); + } + if (padding_ && n_cpf < max_n_cpf) + vdata.insert(vdata.end(), (max_n_cpf - n_cpf) * n_features_cpf_, 0); // Add 0 to unfilled part as padding value + if (max_n_cpf == 0) + vdata.insert(vdata.end(), n_features_cpf_, 0); // Add at least 1 row of 0, otherwise the model has trouble getting output + } + else if (igroup == kNeutralCandidates) { + unsigned int n_npf = features.n_pf_features.size(); + n_npf = std::clamp(n_npf, (unsigned int)0, (unsigned int)25); + + for (unsigned int count = 0; count < n_npf; count++) { + vdata.push_back(features.n_pf_features.at(count).ptrel); + vdata.push_back(features.n_pf_features.at(count).etarel); + vdata.push_back(features.n_pf_features.at(count).phirel); + vdata.push_back(features.n_pf_features.at(count).deltaR); + vdata.push_back(features.n_pf_features.at(count).isGamma); + vdata.push_back(features.n_pf_features.at(count).hadFrac); + vdata.push_back(features.n_pf_features.at(count).drminsv); + vdata.push_back(features.n_pf_features.at(count).puppiw); + } + if (padding_ && n_npf < max_n_npf) + vdata.insert(vdata.end(), (max_n_npf - n_npf) * n_features_npf_, 0); + if (max_n_npf == 0) + vdata.insert(vdata.end(), n_features_npf_, 0); + } + else if (igroup == kVertices) { + unsigned int n_vtx= features.sv_features.size(); + n_vtx = std::clamp(n_vtx, (unsigned int)0, (unsigned int)5); + + for (unsigned int count = 0; count < n_vtx; count++) { + vdata.push_back(features.sv_features.at(count).pt); + vdata.push_back(features.sv_features.at(count).deltaR); + vdata.push_back(features.sv_features.at(count).mass); + vdata.push_back(features.sv_features.at(count).etarel); + vdata.push_back(features.sv_features.at(count).phirel); + vdata.push_back(features.sv_features.at(count).ntracks); + vdata.push_back(features.sv_features.at(count).chi2); + vdata.push_back(features.sv_features.at(count).normchi2); + vdata.push_back(features.sv_features.at(count).dxy); + vdata.push_back(features.sv_features.at(count).dxysig); + vdata.push_back(features.sv_features.at(count).d3d); + vdata.push_back(features.sv_features.at(count).d3dsig); + vdata.push_back(features.sv_features.at(count).costhetasvpv); + vdata.push_back(features.sv_features.at(count).enratio); + } + if (padding_ && n_vtx < max_n_vtx) + vdata.insert(vdata.end(), (max_n_vtx - n_vtx) * n_features_sv_, 0); + if (max_n_vtx == 0) + vdata.insert(vdata.end(), n_features_sv_, 0); + } + else if (igroup == kChargedCandidates4Vec) { + unsigned int n_cpf = features.c_pf_features.size(); + n_cpf = std::clamp(n_cpf, (unsigned int)0, (unsigned int)25); + + for (unsigned int count = 0; count < n_cpf; count++) { + const auto& cpf_pairwise_features = features.c_pf_features.at(count); + vdata.push_back(cpf_pairwise_features.px); + vdata.push_back(cpf_pairwise_features.py); + vdata.push_back(cpf_pairwise_features.pz); + vdata.push_back(cpf_pairwise_features.e); + } + if (padding_ && n_cpf < max_n_cpf) + vdata.insert(vdata.end(), (max_n_cpf - n_cpf) * n_pairwise_features_cpf_, 0); + if (max_n_cpf == 0) + vdata.insert(vdata.end(), n_pairwise_features_cpf_, 0); + } + else if (igroup == kNeutralCandidates4Vec) { + unsigned int n_npf = features.n_pf_features.size(); + n_npf = std::clamp(n_npf, (unsigned int)0, (unsigned int)25); + for (unsigned int count = 0; count < n_npf; count++) { + const auto& npf_pairwise_features = features.n_pf_features.at(count); + vdata.push_back(npf_pairwise_features.px); + vdata.push_back(npf_pairwise_features.py); + vdata.push_back(npf_pairwise_features.pz); + vdata.push_back(npf_pairwise_features.e); + } + if (padding_ && n_npf < max_n_npf) + vdata.insert(vdata.end(), (max_n_npf - n_npf) * n_pairwise_features_npf_, 0); + if (max_n_npf == 0) + vdata.insert(vdata.end(), n_pairwise_features_npf_, 0); + } + else if (igroup == kVertices4Vec ) { + unsigned int n_vtx = features.sv_features.size(); + n_vtx = std::clamp(n_vtx, (unsigned int)0, (unsigned int)5); + for (unsigned int count = 0; count < n_vtx; count++) { + const auto& sv_pairwise_features = features.sv_features.at(count); + vdata.push_back(sv_pairwise_features.px); + vdata.push_back(sv_pairwise_features.py); + vdata.push_back(sv_pairwise_features.pz); + vdata.push_back(sv_pairwise_features.e); + } + if (padding_ && n_vtx < max_n_vtx) + vdata.insert(vdata.end(), (max_n_vtx - n_vtx) * n_pairwise_features_sv_, 0); + if (max_n_vtx == 0) + vdata.insert(vdata.end(), n_pairwise_features_sv_, 0); + } + } + 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..432fd70e3dd09 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"), # 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 From 7fabe9f21c54faed7c2222bd2b9dadecc0b967aa Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Tue, 2 Apr 2024 14:48:41 -0400 Subject: [PATCH 2/9] add particle transformer modofier --- Configuration/ProcessModifiers/python/allSonicTriton_cff.py | 1 + 1 file changed, 1 insertion(+) diff --git a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py index 27291329e6066..f349808c487f7 100644 --- a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py +++ b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py @@ -5,6 +5,7 @@ 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,particleNetPTSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton) From 0301af14c3ed612723d56e7c581cd9e05bd3de48 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Wed, 10 Apr 2024 14:56:51 -0400 Subject: [PATCH 3/9] move common input for particle transformer with/without sonic to interface/ and src/ tensor_filler and tensor_config --- .../ONNXRuntime/interface/tensor_configs.h | 19 ++ .../ONNXRuntime/interface/tensor_fillers.h | 25 +++ ...rticleTransformerAK4ONNXJetTagsProducer.cc | 162 ++++++------------ ...ticleTransformerAK4SonicJetTagsProducer.cc | 143 +++++----------- RecoBTag/ONNXRuntime/src/tensor_fillers.cc | 77 +++++++++ 5 files changed, 211 insertions(+), 215 deletions(-) diff --git a/RecoBTag/ONNXRuntime/interface/tensor_configs.h b/RecoBTag/ONNXRuntime/interface/tensor_configs.h index e528b4e92f74a..a2a1416335375 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_configs.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_configs.h @@ -28,4 +28,23 @@ namespace deepvertex { } // namespace deepvertex +namespace parT { + + enum InputIndexes { + kChargedCandidates = 0, + kNeutralCandidates = 1, + kVertices = 2, + kChargedCandidates4Vec = 3, + kNeutralCandidates4Vec = 4, + kVertices4Vec = 5 + }; + + constexpr unsigned n_features_cpf = 16; + constexpr unsigned n_pairwise_features_cpf = 4; + constexpr unsigned n_features_npf = 8; + constexpr unsigned n_pairwise_features_npf = 4; + constexpr unsigned n_features_sv = 14; + constexpr unsigned n_pairwise_features_sv = 4; +} // namespace parT + #endif diff --git a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h index bb1a25a2d0f0b..6db76cb6e2a41 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h @@ -2,6 +2,7 @@ #define RecoBTag_ONNXRuntime_tensor_fillers_h #include "DataFormats/BTauReco/interface/DeepFlavourTagInfo.h" +#include "RecoBTag/ONNXRuntime/interface/tensor_configs.h" namespace btagbtvdeep { @@ -18,6 +19,30 @@ namespace btagbtvdeep { void seedTrack_tensor_filler(float*& ptr, const btagbtvdeep::SeedingTrackFeatures& seed_features); void neighbourTrack_tensor_filler(float*& ptr, const btagbtvdeep::TrackPairFeatures& neighbourTrack_features); + + std::vector inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputIndexes idx); + + std::vector inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputIndexes idx); + + std::vector inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputIndexes idx); + + template + void parT_tensor_filler(float*& ptr, parT::InputIndexes idx , const parT_features pf) { + std::vector inputs; + inputs = inputs_parT(pf, idx); + for (unsigned int i = 0; i < inputs.size(); i++) { + *ptr = inputs[i]; + ++ptr; + } + if (inputs.size() > 0) --ptr; + } + + template + void parT_tensor_filler(std::vector& vdata, parT::InputIndexes idx , const parT_features pf) { + std::vector inputs; + inputs = inputs_parT(pf, idx); + vdata.insert(vdata.end(), inputs.begin(), inputs.end()); + } } // namespace btagbtvdeep diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc index b27a9b397e520..08da2f566e00a 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_features_cpf}, + {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::n_features_npf}, + {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::n_features_sv}, + {(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::n_pairwise_features_cpf}, + {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::n_pairwise_features_npf}, + {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::n_pairwise_features_sv}}; outputs = globalCache()->run(input_names_, data_, input_shapes_, output_names_, 1)[0]; assert(outputs.size() == flav_names_.size()); @@ -151,24 +139,21 @@ 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_ = features.c_pf_features.size(); + n_npf_ = features.n_pf_features.size(); + n_sv_ = 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(n_cpf_, (unsigned int)1, (unsigned int)25); + n_npf_ = std::clamp(n_npf_, (unsigned int)1, (unsigned int)25); + n_sv_ = std::clamp(n_sv_, (unsigned int)1, (unsigned int)5); - 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_features_cpf, + n_npf_ * parT::n_features_npf, + n_sv_ * parT::n_features_sv, + n_cpf_ * parT::n_pairwise_features_cpf, + n_npf_ * parT::n_pairwise_features_npf, + n_sv_ * parT::n_pairwise_features_sv, }; // init data storage data_.clear(); @@ -185,110 +170,63 @@ void ParticleTransformerAK4ONNXJetTagsProducer::make_inputs(btagbtvdeep::Particl unsigned offset = 0; // c_pf candidates - auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_); + const 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_]; + ptr = &data_[parT::kChargedCandidates][offset + c_pf_n * parT::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); + parT_tensor_filler(ptr, parT::kChargedCandidates, c_pf_features); + assert(start + parT::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_); + const 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_]; + ptr = &data_[parT::kNeutralCandidates][offset + n_pf_n * parT::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); + parT_tensor_filler(ptr, parT::kNeutralCandidates, n_pf_features); + assert(start + parT::n_features_npf - 1 == ptr); } // sv candidates - auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_); + const 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_]; + ptr = &data_[parT::kVertices][offset + sv_n * parT::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); + parT_tensor_filler(ptr, parT::kVertices, sv_features); + assert(start + parT::n_features_sv - 1 == ptr); } // cpf pairwise features (4-vectors) - auto max_cpf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_); + const 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_]; + ptr = &data_[parT::kChargedCandidates4Vec][offset + cpf_n * parT::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(ptr, parT::kChargedCandidates4Vec, cpf_pairwise_features); + assert(start + parT::n_pairwise_features_cpf - 1 == ptr); } // npf pairwise features (4-vectors) - auto max_npf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_); + const 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_]; + ptr = &data_[parT::kNeutralCandidates4Vec][offset + npf_n * parT::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(ptr, parT::kNeutralCandidates4Vec, npf_pairwise_features); + assert(start + parT::n_pairwise_features_npf - 1 == ptr); } // sv pairwise features (4-vectors) - auto max_sv_N = std::min(features.sv_features.size(), (std::size_t)n_sv_); + const 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_]; + ptr = &data_[parT::kVertices4Vec][offset + sv_N * parT::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(ptr, parT::kVertices4Vec, sv_pairwise_features); + assert(start + parT::n_pairwise_features_sv - 1 == ptr); } } diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc index 989ebb95a5071..d78b14fb12aa1 100644 --- a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc +++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc @@ -14,9 +14,11 @@ #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&); @@ -36,21 +38,6 @@ class ParticleTransformerAK4SonicJetTagsProducer : public TritonEDProducer<> { std::vector input_names_; std::vector output_names_; - enum InputIndexes { - kChargedCandidates = 0, - kNeutralCandidates = 1, - kVertices = 2, - kChargedCandidates4Vec = 3, - kNeutralCandidates4Vec = 4, - kVertices4Vec = 5 - }; - - constexpr static unsigned n_features_cpf_ = 16; - constexpr static unsigned n_pairwise_features_cpf_ = 4; - constexpr static unsigned n_features_npf_ = 8; - constexpr static unsigned n_pairwise_features_npf_ = 4; - constexpr static unsigned n_features_sv_ = 14; - constexpr static unsigned n_pairwise_features_sv_ = 4; bool skippedInference_ = false; bool padding_ = true; }; @@ -88,40 +75,40 @@ void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEven client_->setBatchSize(tag_infos->size()); skippedInference_ = false; if (tag_infos->empty()) return; - unsigned int max_n_cpf = 0; - unsigned int max_n_npf = 0; - unsigned int max_n_vtx = 0; + unsigned int max_n_cpf_counter = 0; + unsigned int max_n_npf_counter = 0; + unsigned int max_n_vtx_counter = 0; // Find the max n_cpf, n_npf and n_vtx among all the jets in an event. for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) { - max_n_cpf = std::max(max_n_cpf, + max_n_cpf_counter = std::max(max_n_cpf_counter, static_cast(((*tag_infos)[jet_n]).features().c_pf_features.size())); - max_n_npf = std::max(max_n_npf, + max_n_npf_counter = std::max(max_n_npf_counter, static_cast(((*tag_infos)[jet_n]).features().n_pf_features.size())); - max_n_vtx = std::max(max_n_vtx, + 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 == 0 && max_n_npf == 0 && max_n_vtx == 0) { + 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 - max_n_cpf = std::clamp(max_n_cpf, (unsigned int)0, (unsigned int)25); - max_n_npf = std::clamp(max_n_npf, (unsigned int)0, (unsigned int)25); - max_n_vtx = std::clamp(max_n_vtx, (unsigned int)0, (unsigned int)5); + const unsigned int max_n_cpf = std::clamp(max_n_cpf_counter, (unsigned int)0, (unsigned int)25); + const unsigned int max_n_npf = std::clamp(max_n_npf_counter, (unsigned int)0, (unsigned int)25); + const unsigned int max_n_vtx = std::clamp(max_n_vtx_counter, (unsigned int)0, (unsigned int)5); for (unsigned igroup = 0; igroup < input_names_.size(); ++igroup) { const auto &group_name = input_names_[igroup]; auto &input = iInput.at(group_name); unsigned target = 0; - if (igroup == kChargedCandidates || igroup == kChargedCandidates4Vec) target = std::max((unsigned int)1, max_n_cpf); - else if (igroup == kNeutralCandidates || igroup == kNeutralCandidates4Vec) target = std::max((unsigned int)1, max_n_npf); - else if (igroup == kVertices || igroup == kVertices4Vec) target = std::max((unsigned int)1, max_n_vtx); + if (igroup == parT::kChargedCandidates || igroup == parT::kChargedCandidates4Vec) target = std::max((unsigned int)1, max_n_cpf); + else if (igroup == parT::kNeutralCandidates || igroup == parT::kNeutralCandidates4Vec) target = std::max((unsigned int)1, max_n_npf); + else if (igroup == parT::kVertices || igroup == parT::kVertices4Vec) target = std::max((unsigned int)1, max_n_vtx); input.setShape(0, target); auto tdata = input.allocate(true); @@ -130,121 +117,71 @@ void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEven const auto &features = taginfo.features(); auto &vdata = (*tdata)[jet_n]; // Loop through the n cpf, and in the case that n_cpf is smaller than max_n_cpf, add padding values to all features - if (igroup == kChargedCandidates) { + if (igroup == parT::kChargedCandidates) { unsigned int n_cpf = features.c_pf_features.size(); n_cpf = std::clamp(n_cpf, (unsigned int)0, (unsigned int)25); for (unsigned int count = 0; count < n_cpf; count++) { - vdata.push_back(features.c_pf_features.at(count).btagPf_trackEtaRel); - vdata.push_back(features.c_pf_features.at(count).btagPf_trackPtRel); - vdata.push_back(features.c_pf_features.at(count).btagPf_trackPPar); - vdata.push_back(features.c_pf_features.at(count).btagPf_trackDeltaR); - vdata.push_back(features.c_pf_features.at(count).btagPf_trackPParRatio); - vdata.push_back(features.c_pf_features.at(count).btagPf_trackSip2dVal); - vdata.push_back(features.c_pf_features.at(count).btagPf_trackSip2dSig); - vdata.push_back(features.c_pf_features.at(count).btagPf_trackSip3dVal); - vdata.push_back(features.c_pf_features.at(count).btagPf_trackSip3dSig); - vdata.push_back(features.c_pf_features.at(count).btagPf_trackJetDistVal); - vdata.push_back(features.c_pf_features.at(count).ptrel); - vdata.push_back(features.c_pf_features.at(count).drminsv); - vdata.push_back(features.c_pf_features.at(count).vtx_ass); - vdata.push_back(features.c_pf_features.at(count).puppiw); - vdata.push_back(features.c_pf_features.at(count).chi2); - vdata.push_back(features.c_pf_features.at(count).quality); + parT_tensor_filler(vdata, parT::kChargedCandidates, features.c_pf_features.at(count)); } if (padding_ && n_cpf < max_n_cpf) - vdata.insert(vdata.end(), (max_n_cpf - n_cpf) * n_features_cpf_, 0); // Add 0 to unfilled part as padding value + vdata.insert(vdata.end(), (max_n_cpf - n_cpf) * parT::n_features_cpf, 0); // Add 0 to unfilled part as padding value if (max_n_cpf == 0) - vdata.insert(vdata.end(), n_features_cpf_, 0); // Add at least 1 row of 0, otherwise the model has trouble getting output + vdata.insert(vdata.end(), parT::n_features_cpf, 0); // Add at least 1 row of 0 for a jet that has 0 cpf/npf/sv. } - else if (igroup == kNeutralCandidates) { + else if (igroup == parT::kNeutralCandidates) { unsigned int n_npf = features.n_pf_features.size(); n_npf = std::clamp(n_npf, (unsigned int)0, (unsigned int)25); - for (unsigned int count = 0; count < n_npf; count++) { - vdata.push_back(features.n_pf_features.at(count).ptrel); - vdata.push_back(features.n_pf_features.at(count).etarel); - vdata.push_back(features.n_pf_features.at(count).phirel); - vdata.push_back(features.n_pf_features.at(count).deltaR); - vdata.push_back(features.n_pf_features.at(count).isGamma); - vdata.push_back(features.n_pf_features.at(count).hadFrac); - vdata.push_back(features.n_pf_features.at(count).drminsv); - vdata.push_back(features.n_pf_features.at(count).puppiw); + parT_tensor_filler(vdata, parT::kNeutralCandidates, features.n_pf_features.at(count)); } if (padding_ && n_npf < max_n_npf) - vdata.insert(vdata.end(), (max_n_npf - n_npf) * n_features_npf_, 0); + vdata.insert(vdata.end(), (max_n_npf - n_npf) * parT::n_features_npf, 0); // Add 0 to unfilled part as padding value if (max_n_npf == 0) - vdata.insert(vdata.end(), n_features_npf_, 0); + vdata.insert(vdata.end(), parT::n_features_npf, 0); } - else if (igroup == kVertices) { + else if (igroup == parT::kVertices) { unsigned int n_vtx= features.sv_features.size(); n_vtx = std::clamp(n_vtx, (unsigned int)0, (unsigned int)5); - for (unsigned int count = 0; count < n_vtx; count++) { - vdata.push_back(features.sv_features.at(count).pt); - vdata.push_back(features.sv_features.at(count).deltaR); - vdata.push_back(features.sv_features.at(count).mass); - vdata.push_back(features.sv_features.at(count).etarel); - vdata.push_back(features.sv_features.at(count).phirel); - vdata.push_back(features.sv_features.at(count).ntracks); - vdata.push_back(features.sv_features.at(count).chi2); - vdata.push_back(features.sv_features.at(count).normchi2); - vdata.push_back(features.sv_features.at(count).dxy); - vdata.push_back(features.sv_features.at(count).dxysig); - vdata.push_back(features.sv_features.at(count).d3d); - vdata.push_back(features.sv_features.at(count).d3dsig); - vdata.push_back(features.sv_features.at(count).costhetasvpv); - vdata.push_back(features.sv_features.at(count).enratio); + parT_tensor_filler(vdata, parT::kVertices, features.sv_features.at(count)); } if (padding_ && n_vtx < max_n_vtx) - vdata.insert(vdata.end(), (max_n_vtx - n_vtx) * n_features_sv_, 0); + vdata.insert(vdata.end(), (max_n_vtx - n_vtx) * parT::n_features_sv, 0); if (max_n_vtx == 0) - vdata.insert(vdata.end(), n_features_sv_, 0); + vdata.insert(vdata.end(), parT::n_features_sv, 0); } - else if (igroup == kChargedCandidates4Vec) { + else if (igroup == parT::kChargedCandidates4Vec) { unsigned int n_cpf = features.c_pf_features.size(); n_cpf = std::clamp(n_cpf, (unsigned int)0, (unsigned int)25); - for (unsigned int count = 0; count < n_cpf; count++) { - const auto& cpf_pairwise_features = features.c_pf_features.at(count); - vdata.push_back(cpf_pairwise_features.px); - vdata.push_back(cpf_pairwise_features.py); - vdata.push_back(cpf_pairwise_features.pz); - vdata.push_back(cpf_pairwise_features.e); + parT_tensor_filler(vdata, parT::kChargedCandidates4Vec, features.c_pf_features.at(count)); } if (padding_ && n_cpf < max_n_cpf) - vdata.insert(vdata.end(), (max_n_cpf - n_cpf) * n_pairwise_features_cpf_, 0); + vdata.insert(vdata.end(), (max_n_cpf - n_cpf) * parT::n_pairwise_features_cpf, 0); if (max_n_cpf == 0) - vdata.insert(vdata.end(), n_pairwise_features_cpf_, 0); + vdata.insert(vdata.end(), parT::n_pairwise_features_cpf, 0); } - else if (igroup == kNeutralCandidates4Vec) { + else if (igroup == parT::kNeutralCandidates4Vec) { unsigned int n_npf = features.n_pf_features.size(); n_npf = std::clamp(n_npf, (unsigned int)0, (unsigned int)25); for (unsigned int count = 0; count < n_npf; count++) { - const auto& npf_pairwise_features = features.n_pf_features.at(count); - vdata.push_back(npf_pairwise_features.px); - vdata.push_back(npf_pairwise_features.py); - vdata.push_back(npf_pairwise_features.pz); - vdata.push_back(npf_pairwise_features.e); + parT_tensor_filler(vdata, parT::kNeutralCandidates4Vec, features.n_pf_features.at(count)); } if (padding_ && n_npf < max_n_npf) - vdata.insert(vdata.end(), (max_n_npf - n_npf) * n_pairwise_features_npf_, 0); + vdata.insert(vdata.end(), (max_n_npf - n_npf) * parT::n_pairwise_features_npf, 0); if (max_n_npf == 0) - vdata.insert(vdata.end(), n_pairwise_features_npf_, 0); + vdata.insert(vdata.end(), parT::n_pairwise_features_npf, 0); } - else if (igroup == kVertices4Vec ) { + else if (igroup == parT::kVertices4Vec) { unsigned int n_vtx = features.sv_features.size(); n_vtx = std::clamp(n_vtx, (unsigned int)0, (unsigned int)5); for (unsigned int count = 0; count < n_vtx; count++) { - const auto& sv_pairwise_features = features.sv_features.at(count); - vdata.push_back(sv_pairwise_features.px); - vdata.push_back(sv_pairwise_features.py); - vdata.push_back(sv_pairwise_features.pz); - vdata.push_back(sv_pairwise_features.e); + parT_tensor_filler(vdata, parT::kVertices4Vec, features.sv_features.at(count)); } if (padding_ && n_vtx < max_n_vtx) - vdata.insert(vdata.end(), (max_n_vtx - n_vtx) * n_pairwise_features_sv_, 0); + vdata.insert(vdata.end(), (max_n_vtx - n_vtx) * parT::n_pairwise_features_sv, 0); if (max_n_vtx == 0) - vdata.insert(vdata.end(), n_pairwise_features_sv_, 0); + vdata.insert(vdata.end(), parT::n_pairwise_features_sv, 0); } } input.toServer(tdata); diff --git a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc index da49b1b0119d7..81abc653cc8e9 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::InputIndexes idx) { + std::vector inputs; + if (idx == 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 (idx == 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::InputIndexes idx) { + std::vector inputs; + if (idx == 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 (idx == 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::InputIndexes idx) { + std::vector inputs; + if (idx == 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 (idx == 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 From a5b07c83ece0e5573aa82fda61d363f3a4eea656 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Mon, 29 Apr 2024 14:30:13 -0400 Subject: [PATCH 4/9] respond to code comments --- RecoBTag/ONNXRuntime/BuildFile.xml | 1 + .../ONNXRuntime/interface/tensor_configs.h | 52 ++++++--- .../ONNXRuntime/interface/tensor_fillers.h | 43 ++++--- ...rticleTransformerAK4ONNXJetTagsProducer.cc | 99 +++++----------- ...ticleTransformerAK4SonicJetTagsProducer.cc | 106 +++++------------- RecoBTag/ONNXRuntime/src/tensor_fillers.cc | 18 +-- 6 files changed, 128 insertions(+), 191 deletions(-) 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 a2a1416335375..0fbdbb5699698 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; @@ -30,21 +31,46 @@ namespace deepvertex { namespace parT { - enum InputIndexes { - kChargedCandidates = 0, - kNeutralCandidates = 1, - kVertices = 2, - kChargedCandidates4Vec = 3, - kNeutralCandidates4Vec = 4, - kVertices4Vec = 5 + enum InputFeatures { + kChargedCandidates=0, + kNeutralCandidates=1, + kVertices=2, + kChargedCandidates4Vec=3, + kNeutralCandidates4Vec=4, + kVertices4Vec=5 + }; + + const std::map InputIndexes{ + {0, kChargedCandidates}, + {1, kNeutralCandidates}, + {2, kVertices}, + {3, kChargedCandidates4Vec}, + {4, kNeutralCandidates4Vec}, + {5, kVertices4Vec} + }; + + 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} }; - constexpr unsigned n_features_cpf = 16; - constexpr unsigned n_pairwise_features_cpf = 4; - constexpr unsigned n_features_npf = 8; - constexpr unsigned n_pairwise_features_npf = 4; - constexpr unsigned n_features_sv = 14; - constexpr unsigned n_pairwise_features_sv = 4; } // namespace parT #endif diff --git a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h index 6db76cb6e2a41..ef1cc7e69543c 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h @@ -2,6 +2,8 @@ #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 { @@ -20,28 +22,39 @@ namespace btagbtvdeep { void neighbourTrack_tensor_filler(float*& ptr, const btagbtvdeep::TrackPairFeatures& neighbourTrack_features); - std::vector inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputIndexes idx); + std::vector inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputFeatures ifeature); - std::vector inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputIndexes idx); + std::vector inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputFeatures ifeature); - std::vector inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputIndexes idx); + std::vector inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature); template - void parT_tensor_filler(float*& ptr, parT::InputIndexes idx , const parT_features pf) { - std::vector inputs; - inputs = inputs_parT(pf, idx); - for (unsigned int i = 0; i < inputs.size(); i++) { - *ptr = inputs[i]; - ++ptr; + 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); } - if (inputs.size() > 0) --ptr; } - + template - void parT_tensor_filler(std::vector& vdata, parT::InputIndexes idx , const parT_features pf) { - std::vector inputs; - inputs = inputs_parT(pf, idx); - vdata.insert(vdata.end(), inputs.begin(), inputs.end()); + 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 diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc index 08da2f566e00a..810ab82ff6f2a 100644 --- a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc +++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc @@ -112,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)parT::n_features_cpf}, - {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::n_features_npf}, - {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::n_features_sv}, - {(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::n_pairwise_features_cpf}, - {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::n_pairwise_features_npf}, - {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::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()); @@ -139,21 +139,17 @@ void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes( const reco::FeaturesTagInfo taginfo) { const auto& features = taginfo.features(); - n_cpf_ = features.c_pf_features.size(); - n_npf_ = features.n_pf_features.size(); - n_sv_ = features.sv_features.size(); - - n_cpf_ = std::clamp(n_cpf_, (unsigned int)1, (unsigned int)25); - n_npf_ = std::clamp(n_npf_, (unsigned int)1, (unsigned int)25); - n_sv_ = std::clamp(n_sv_, (unsigned int)1, (unsigned int)5); + 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); input_sizes_ = { - n_cpf_ * parT::n_features_cpf, - n_npf_ * parT::n_features_npf, - n_sv_ * parT::n_features_sv, - n_cpf_ * parT::n_pairwise_features_cpf, - n_npf_ * parT::n_pairwise_features_npf, - n_sv_ * parT::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(); @@ -165,69 +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 - const 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_[parT::kChargedCandidates][offset + c_pf_n * parT::n_features_cpf]; - start = ptr; - parT_tensor_filler(ptr, parT::kChargedCandidates, c_pf_features); - assert(start + parT::n_features_cpf - 1 == ptr); - } + auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_); + auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_); + auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_); + // c_pf candidates + parT_tensor_filler(data_, parT::kChargedCandidates, features.c_pf_features, max_c_pf_n, start, offset); // n_pf candidates - const 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_[parT::kNeutralCandidates][offset + n_pf_n * parT::n_features_npf]; - start = ptr; - parT_tensor_filler(ptr, parT::kNeutralCandidates, n_pf_features); - assert(start + parT::n_features_npf - 1 == ptr); - } - + parT_tensor_filler(data_, parT::kNeutralCandidates, features.n_pf_features, max_n_pf_n, start, offset); // sv candidates - const 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_[parT::kVertices][offset + sv_n * parT::n_features_sv]; - start = ptr; - parT_tensor_filler(ptr, parT::kVertices, sv_features); - assert(start + parT::n_features_sv - 1 == ptr); - } - + parT_tensor_filler(data_, parT::kVertices, features.sv_features, max_sv_n, start, offset); // cpf pairwise features (4-vectors) - const 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_[parT::kChargedCandidates4Vec][offset + cpf_n * parT::n_pairwise_features_cpf]; - start = ptr; - parT_tensor_filler(ptr, parT::kChargedCandidates4Vec, cpf_pairwise_features); - assert(start + parT::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) - const 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_[parT::kNeutralCandidates4Vec][offset + npf_n * parT::n_pairwise_features_npf]; - start = ptr; - parT_tensor_filler(ptr, parT::kNeutralCandidates4Vec, npf_pairwise_features); - assert(start + parT::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) - const 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_[parT::kVertices4Vec][offset + sv_N * parT::n_pairwise_features_sv]; - start = ptr; - parT_tensor_filler(ptr, parT::kVertices4Vec, sv_pairwise_features); - assert(start + parT::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 index d78b14fb12aa1..346d1c5fb2577 100644 --- a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc +++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc @@ -39,7 +39,6 @@ class ParticleTransformerAK4SonicJetTagsProducer : public TritonEDProducer<> { std::vector output_names_; bool skippedInference_ = false; - bool padding_ = true; }; ParticleTransformerAK4SonicJetTagsProducer::ParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet& iConfig) @@ -75,11 +74,11 @@ void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEven 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; - - // Find the max n_cpf, n_npf and n_vtx among all the jets in an event. 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())); @@ -97,92 +96,37 @@ void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEven } // 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 max_n_cpf = std::clamp(max_n_cpf_counter, (unsigned int)0, (unsigned int)25); - const unsigned int max_n_npf = std::clamp(max_n_npf_counter, (unsigned int)0, (unsigned int)25); - const unsigned int max_n_vtx = std::clamp(max_n_vtx_counter, (unsigned int)0, (unsigned int)5); - + 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} + }; + for (unsigned igroup = 0; igroup < input_names_.size(); ++igroup) { const auto &group_name = input_names_[igroup]; auto &input = iInput.at(group_name); - unsigned target = 0; - - if (igroup == parT::kChargedCandidates || igroup == parT::kChargedCandidates4Vec) target = std::max((unsigned int)1, max_n_cpf); - else if (igroup == parT::kNeutralCandidates || igroup == parT::kNeutralCandidates4Vec) target = std::max((unsigned int)1, max_n_npf); - else if (igroup == parT::kVertices || igroup == parT::kVertices4Vec) target = std::max((unsigned int)1, max_n_vtx); - - input.setShape(0, target); + + const parT::InputFeatures ifeature = parT::InputIndexes.at(igroup); + 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]; - // Loop through the n cpf, and in the case that n_cpf is smaller than max_n_cpf, add padding values to all features - if (igroup == parT::kChargedCandidates) { - unsigned int n_cpf = features.c_pf_features.size(); - n_cpf = std::clamp(n_cpf, (unsigned int)0, (unsigned int)25); - for (unsigned int count = 0; count < n_cpf; count++) { - parT_tensor_filler(vdata, parT::kChargedCandidates, features.c_pf_features.at(count)); - } - if (padding_ && n_cpf < max_n_cpf) - vdata.insert(vdata.end(), (max_n_cpf - n_cpf) * parT::n_features_cpf, 0); // Add 0 to unfilled part as padding value - if (max_n_cpf == 0) - vdata.insert(vdata.end(), parT::n_features_cpf, 0); // Add at least 1 row of 0 for a jet that has 0 cpf/npf/sv. - } - else if (igroup == parT::kNeutralCandidates) { - unsigned int n_npf = features.n_pf_features.size(); - n_npf = std::clamp(n_npf, (unsigned int)0, (unsigned int)25); - for (unsigned int count = 0; count < n_npf; count++) { - parT_tensor_filler(vdata, parT::kNeutralCandidates, features.n_pf_features.at(count)); - } - if (padding_ && n_npf < max_n_npf) - vdata.insert(vdata.end(), (max_n_npf - n_npf) * parT::n_features_npf, 0); // Add 0 to unfilled part as padding value - if (max_n_npf == 0) - vdata.insert(vdata.end(), parT::n_features_npf, 0); - } - else if (igroup == parT::kVertices) { - unsigned int n_vtx= features.sv_features.size(); - n_vtx = std::clamp(n_vtx, (unsigned int)0, (unsigned int)5); - for (unsigned int count = 0; count < n_vtx; count++) { - parT_tensor_filler(vdata, parT::kVertices, features.sv_features.at(count)); - } - if (padding_ && n_vtx < max_n_vtx) - vdata.insert(vdata.end(), (max_n_vtx - n_vtx) * parT::n_features_sv, 0); - if (max_n_vtx == 0) - vdata.insert(vdata.end(), parT::n_features_sv, 0); - } - else if (igroup == parT::kChargedCandidates4Vec) { - unsigned int n_cpf = features.c_pf_features.size(); - n_cpf = std::clamp(n_cpf, (unsigned int)0, (unsigned int)25); - for (unsigned int count = 0; count < n_cpf; count++) { - parT_tensor_filler(vdata, parT::kChargedCandidates4Vec, features.c_pf_features.at(count)); - } - if (padding_ && n_cpf < max_n_cpf) - vdata.insert(vdata.end(), (max_n_cpf - n_cpf) * parT::n_pairwise_features_cpf, 0); - if (max_n_cpf == 0) - vdata.insert(vdata.end(), parT::n_pairwise_features_cpf, 0); - } - else if (igroup == parT::kNeutralCandidates4Vec) { - unsigned int n_npf = features.n_pf_features.size(); - n_npf = std::clamp(n_npf, (unsigned int)0, (unsigned int)25); - for (unsigned int count = 0; count < n_npf; count++) { - parT_tensor_filler(vdata, parT::kNeutralCandidates4Vec, features.n_pf_features.at(count)); - } - if (padding_ && n_npf < max_n_npf) - vdata.insert(vdata.end(), (max_n_npf - n_npf) * parT::n_pairwise_features_npf, 0); - if (max_n_npf == 0) - vdata.insert(vdata.end(), parT::n_pairwise_features_npf, 0); - } - else if (igroup == parT::kVertices4Vec) { - unsigned int n_vtx = features.sv_features.size(); - n_vtx = std::clamp(n_vtx, (unsigned int)0, (unsigned int)5); - for (unsigned int count = 0; count < n_vtx; count++) { - parT_tensor_filler(vdata, parT::kVertices4Vec, features.sv_features.at(count)); - } - if (padding_ && n_vtx < max_n_vtx) - vdata.insert(vdata.end(), (max_n_vtx - n_vtx) * parT::n_pairwise_features_sv, 0); - if (max_n_vtx == 0) - vdata.insert(vdata.end(), parT::n_pairwise_features_sv, 0); - } + 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); } diff --git a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc index 81abc653cc8e9..309d4a1e4b55b 100644 --- a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc +++ b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc @@ -138,9 +138,9 @@ namespace btagbtvdeep { *(++ptr) = neighbourTrack_features.dphi_PCAjetDirs; } - std::vector inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputIndexes idx) { + std::vector inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputFeatures ifeature) { std::vector inputs; - if (idx == parT::kChargedCandidates) { + 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); @@ -158,7 +158,7 @@ namespace btagbtvdeep { inputs.push_back(c_pf_features.chi2); inputs.push_back(c_pf_features.quality); } - else if (idx == parT::kChargedCandidates4Vec) { + 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); @@ -167,9 +167,9 @@ namespace btagbtvdeep { return inputs; } - std::vector inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputIndexes idx) { + std::vector inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputFeatures ifeature) { std::vector inputs; - if (idx == parT::kNeutralCandidates) { + 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); @@ -179,7 +179,7 @@ namespace btagbtvdeep { inputs.push_back(n_pf_features.drminsv); inputs.push_back(n_pf_features.puppiw); } - else if (idx == parT::kNeutralCandidates4Vec) { + 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); @@ -188,9 +188,9 @@ namespace btagbtvdeep { return inputs; } - std::vector inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputIndexes idx) { + std::vector inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature) { std::vector inputs; - if (idx == parT::kVertices) { + if (ifeature == parT::kVertices) { inputs.push_back(sv_features.pt); inputs.push_back(sv_features.deltaR); inputs.push_back(sv_features.mass); @@ -206,7 +206,7 @@ namespace btagbtvdeep { inputs.push_back(sv_features.costhetasvpv); inputs.push_back(sv_features.enratio); } - else if (idx == parT::kVertices4Vec) { + else if (ifeature == parT::kVertices4Vec) { inputs.push_back(sv_features.px); inputs.push_back(sv_features.py); inputs.push_back(sv_features.pz); From 35d7a1824cdc58f75e322f51ee8ca6ffd5193449 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Fri, 3 May 2024 10:44:59 -0400 Subject: [PATCH 5/9] remove InputIndexes and loop through enum --- RecoBTag/ONNXRuntime/interface/tensor_configs.h | 17 +++++------------ ...articleTransformerAK4SonicJetTagsProducer.cc | 12 ++++++------ 2 files changed, 11 insertions(+), 18 deletions(-) diff --git a/RecoBTag/ONNXRuntime/interface/tensor_configs.h b/RecoBTag/ONNXRuntime/interface/tensor_configs.h index 0fbdbb5699698..e52a94b29ae6e 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_configs.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_configs.h @@ -32,23 +32,16 @@ namespace deepvertex { namespace parT { enum InputFeatures { - kChargedCandidates=0, + kBegin=0, + kChargedCandidates=kBegin, kNeutralCandidates=1, kVertices=2, kChargedCandidates4Vec=3, kNeutralCandidates4Vec=4, - kVertices4Vec=5 + kVertices4Vec=5, + kEnd=6 }; - - const std::map InputIndexes{ - {0, kChargedCandidates}, - {1, kNeutralCandidates}, - {2, kVertices}, - {3, kChargedCandidates4Vec}, - {4, kNeutralCandidates4Vec}, - {5, kVertices4Vec} - }; - + constexpr unsigned n_cpf_accept = 25; constexpr unsigned n_npf_accept = 25; constexpr unsigned n_sv_accept = 5; diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc index 346d1c5fb2577..cf515d170373e 100644 --- a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc +++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc @@ -108,12 +108,12 @@ void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEven {parT::kNeutralCandidates4Vec, target_n_npf}, {parT::kVertices4Vec, target_n_vtx} }; - - for (unsigned igroup = 0; igroup < input_names_.size(); ++igroup) { - const auto &group_name = input_names_[igroup]; + + // 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); - - const parT::InputFeatures ifeature = parT::InputIndexes.at(igroup); + input.setShape(0, target_n.at(ifeature)); auto tdata = input.allocate(true); @@ -129,7 +129,7 @@ void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEven 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) { From 2b6404d3e155dd1b27fcda6d7928ea3bcebf2361 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Tue, 6 Aug 2024 09:50:59 -0400 Subject: [PATCH 6/9] code patch --- .../ONNXRuntime/interface/tensor_configs.h | 56 +++++++-------- .../ONNXRuntime/interface/tensor_fillers.h | 34 +++++++--- ...rticleTransformerAK4ONNXJetTagsProducer.cc | 2 +- ...ticleTransformerAK4SonicJetTagsProducer.cc | 68 ++++++++++--------- RecoBTag/ONNXRuntime/src/tensor_fillers.cc | 22 +++--- 5 files changed, 97 insertions(+), 85 deletions(-) diff --git a/RecoBTag/ONNXRuntime/interface/tensor_configs.h b/RecoBTag/ONNXRuntime/interface/tensor_configs.h index e52a94b29ae6e..085e179863058 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_configs.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_configs.h @@ -32,38 +32,34 @@ namespace deepvertex { namespace parT { enum InputFeatures { - kBegin=0, - kChargedCandidates=kBegin, - kNeutralCandidates=1, - kVertices=2, - kChargedCandidates4Vec=3, - kNeutralCandidates4Vec=4, - kVertices4Vec=5, - kEnd=6 + 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 + 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 #endif diff --git a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h index ef1cc7e69543c..db3facc2338d1 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_fillers.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_fillers.h @@ -21,15 +21,22 @@ namespace btagbtvdeep { void seedTrack_tensor_filler(float*& ptr, const btagbtvdeep::SeedingTrackFeatures& seed_features); 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::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) { + 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); @@ -40,21 +47,26 @@ namespace btagbtvdeep { *ptr = inputs[i]; ++ptr; } - if (inputs.size() > 0) --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)); + + 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 + vdata.insert(vdata.end(), (target_n - n) * n_features, 0); // Add 0 to unfilled part as padding value } } // namespace btagbtvdeep diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc index 810ab82ff6f2a..ea5fc6b0020f4 100644 --- a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc +++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc @@ -30,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; diff --git a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc index cf515d170373e..2a33e976e61f0 100644 --- a/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc +++ b/RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc @@ -21,13 +21,13 @@ class ParticleTransformerAK4SonicJetTagsProducer : public TritonEDProducer<> { public: - explicit ParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet&); + 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&); + static void fillDescriptions(edm::ConfigurationDescriptions &); private: typedef std::vector TagInfoCollection; @@ -41,21 +41,21 @@ class ParticleTransformerAK4SonicJetTagsProducer : public TritonEDProducer<> { bool skippedInference_ = false; }; -ParticleTransformerAK4SonicJetTagsProducer::ParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet& iConfig) +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_) { + for (const auto &flav_name : flav_names_) { produces(flav_name); } } ParticleTransformerAK4SonicJetTagsProducer::~ParticleTransformerAK4SonicJetTagsProducer() {} -void ParticleTransformerAK4SonicJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) { +void ParticleTransformerAK4SonicJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions &descriptions) { // pfParticleTransformerAK4JetTags edm::ParameterSetDescription desc; TritonClient::fillPSetDescription(desc); @@ -68,49 +68,51 @@ void ParticleTransformerAK4SonicJetTagsProducer::fillDescriptions(edm::Configura descriptions.add("pfParticleTransformerAK4SonicJetTags", desc); } -void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) { +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; + 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; + // 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())); + 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} - }; - + 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)) { + 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); @@ -121,9 +123,9 @@ void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEven const auto &taginfo = (*tag_infos)[jet_n]; const auto &features = taginfo.features(); auto &vdata = (*tdata)[jet_n]; - if (ifeature == parT::kChargedCandidates || ifeature == parT::kChargedCandidates4Vec) + 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) + 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); @@ -132,7 +134,9 @@ void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEven } } -void ParticleTransformerAK4SonicJetTagsProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup, Output const &iOutput) { +void ParticleTransformerAK4SonicJetTagsProducer::produce(edm::Event &iEvent, + const edm::EventSetup &iSetup, + Output const &iOutput) { edm::Handle tag_infos; iEvent.getByToken(src_, tag_infos); diff --git a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc index 309d4a1e4b55b..9ecdf0302b45a 100644 --- a/RecoBTag/ONNXRuntime/src/tensor_fillers.cc +++ b/RecoBTag/ONNXRuntime/src/tensor_fillers.cc @@ -138,7 +138,8 @@ namespace btagbtvdeep { *(++ptr) = neighbourTrack_features.dphi_PCAjetDirs; } - std::vector inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features, parT::InputFeatures ifeature) { + 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); @@ -157,8 +158,7 @@ namespace btagbtvdeep { 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) { + } 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); @@ -167,7 +167,8 @@ namespace btagbtvdeep { return inputs; } - std::vector inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features, parT::InputFeatures ifeature) { + 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); @@ -178,17 +179,17 @@ namespace btagbtvdeep { 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) { + } 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); + inputs.push_back(n_pf_features.e); } return inputs; } - std::vector inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature) { + 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); @@ -205,12 +206,11 @@ namespace btagbtvdeep { inputs.push_back(sv_features.d3dsig); inputs.push_back(sv_features.costhetasvpv); inputs.push_back(sv_features.enratio); - } - else if (ifeature == parT::kVertices4Vec) { + } 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); + inputs.push_back(sv_features.e); } return inputs; } From 28d81de515baca003ae0a30e15ec9b0b8e6091ea Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Tue, 6 Aug 2024 11:41:18 -0400 Subject: [PATCH 7/9] remove unnecessary comments --- RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py b/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py index 432fd70e3dd09..4a7fd51782280 100644 --- a/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py +++ b/RecoBTag/ONNXRuntime/python/pfParticleTransformerAK4_cff.py @@ -14,8 +14,8 @@ 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 + 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), From 7a4032ed957b098bf86a5afb219d9df41c24fbfb Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Mon, 12 Aug 2024 10:37:32 -0400 Subject: [PATCH 8/9] resolve comments to make variables inline and use array instead of map --- .../ONNXRuntime/interface/tensor_configs.h | 60 ++++++++++--------- 1 file changed, 31 insertions(+), 29 deletions(-) diff --git a/RecoBTag/ONNXRuntime/interface/tensor_configs.h b/RecoBTag/ONNXRuntime/interface/tensor_configs.h index 085e179863058..f0f49a8b5cd5c 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_configs.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_configs.h @@ -1,7 +1,7 @@ #ifndef RecoBTag_ONNXRuntime_tensor_configs_h #define RecoBTag_ONNXRuntime_tensor_configs_h -#include +#include namespace deepflavour { constexpr unsigned n_features_global = 15; @@ -32,34 +32,36 @@ namespace deepvertex { namespace parT { enum InputFeatures { - kBegin = 0, - kChargedCandidates = kBegin, - kNeutralCandidates = 1, - kVertices = 2, - kChargedCandidates4Vec = 3, - kNeutralCandidates4Vec = 4, - kVertices4Vec = 5, - kEnd = 6 + 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 + + 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 From df4fe14fd82b9f867c5149f60389ec45e7a581e0 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Mon, 12 Aug 2024 10:54:30 -0400 Subject: [PATCH 9/9] update format --- .../ONNXRuntime/interface/tensor_configs.h | 50 ++++++++++--------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/RecoBTag/ONNXRuntime/interface/tensor_configs.h b/RecoBTag/ONNXRuntime/interface/tensor_configs.h index f0f49a8b5cd5c..98faeda8403db 100644 --- a/RecoBTag/ONNXRuntime/interface/tensor_configs.h +++ b/RecoBTag/ONNXRuntime/interface/tensor_configs.h @@ -32,36 +32,38 @@ namespace deepvertex { namespace parT { enum InputFeatures { - kBegin=0, - kChargedCandidates=kBegin, - kNeutralCandidates=1, - kVertices=2, - kChargedCandidates4Vec=3, - kNeutralCandidates4Vec=4, - kVertices4Vec=5, - kEnd=6 + 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 + 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 + 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 +} // namespace parT #endif