Particle Transformer sonic producer #45579

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

cmsbuild merged 9 commits into cms-sw:master from y19y19:CMSSW_14_1_0_pre0_testParT

Aug 20, 2024

Configuration/ProcessModifiers/python/allSonicTriton_cff.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -5,6 +5,8 @@ @@
     from Configuration.ProcessModifiers.particleNetPTSonicTriton_cff import particleNetPTSonicTriton
     from Configuration.ProcessModifiers.deepMETSonicTriton_cff import deepMETSonicTriton
     from Configuration.ProcessModifiers.deepTauSonicTriton_cff import deepTauSonicTriton
+    from Configuration.ProcessModifiers.particleTransformerAK4SonicTriton_cff import particleTransformerAK4SonicTriton
     # collect all SonicTriton-related process modifiers here
-    allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,deepTauSonicTriton)
+    allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,particleNetPTSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton)

Configuration/ProcessModifiers/python/particleTransformerAK4SonicTriton_cff.py

Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		import FWCore.ParameterSet.Config as cms

		particleTransformerAK4SonicTriton = cms.Modifier()

RecoBTag/ONNXRuntime/BuildFile.xml

-Original file line number
+Diff line change
@@ -1,3 +1,4 @@
+    <use name="PhysicsTools/ONNXRuntime"/>
     <use name="DataFormats/BTauReco"/>
     <export>
       <lib name="1"/>
@@ Expand Down @@

RecoBTag/ONNXRuntime/interface/tensor_configs.h

-Original file line number
+Diff line change
@@ -1,6 +1,7 @@
     #ifndef RecoBTag_ONNXRuntime_tensor_configs_h
     #define RecoBTag_ONNXRuntime_tensor_configs_h
+    #include <array>
     namespace deepflavour {
       constexpr unsigned n_features_global = 15;
@@ Expand Down Expand Up / @@ -28,4 +29,41 @@ namespace deepvertex { @@
     }  // namespace deepvertex
+    namespace parT {
+      enum InputFeatures {
+        kBegin = 0,
+        kChargedCandidates = kBegin,
+        kNeutralCandidates = 1,
+        kVertices = 2,
+        kChargedCandidates4Vec = 3,
+        kNeutralCandidates4Vec = 4,
+        kVertices4Vec = 5,
+        kEnd = 6
+      };
+      inline constexpr unsigned n_cpf_accept = 25;
+      inline constexpr unsigned n_npf_accept = 25;
+      inline constexpr unsigned n_sv_accept = 5;
+      constexpr std::array<unsigned int, kEnd> N_InputFeatures{{
+,  // kChargedCandidates
+,   // kNeutralCandidates
+,  // kVertices
+,   // kChargedCandidates4Vec
+,   // kNeutralCandidates4Vec
+,   // kVertices4Vec
+      }};
+      constexpr std::array<unsigned int, kEnd> 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

RecoBTag/ONNXRuntime/interface/tensor_fillers.h

-Original file line number
+Diff line change
@@ Expand Up / @@ -2,6 +2,9 @@ @@
     #define RecoBTag_ONNXRuntime_tensor_fillers_h
     #include "DataFormats/BTauReco/interface/DeepFlavourTagInfo.h"
+    #include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h"
+    #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
+    #include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
     namespace btagbtvdeep {
@@ Expand All / @@ -19,6 +22,53 @@ namespace btagbtvdeep { @@
       void neighbourTrack_tensor_filler(float*& ptr, const btagbtvdeep::TrackPairFeatures& neighbourTrack_features);
+      std::vector<float> inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features,
+                                     parT::InputFeatures ifeature);
+      std::vector<float> inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features,
+                                     parT::InputFeatures ifeature);
+      std::vector<float> inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature);
+      template <class parT_features>
+      void parT_tensor_filler(cms::Ort::FloatArrays& data,
+                              const parT::InputFeatures ifeature,
+                              const std::vector<parT_features>& 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<float>& inputs = inputs_parT(f, ifeature);
+          for (unsigned int i = 0; i < inputs.size(); i++) {
+            *ptr = inputs[i];
+            ++ptr;
+          }
+          if (inputs.size() > 0)
+            --ptr;
+          assert(start + parT::N_InputFeatures.at(ifeature) - 1 == ptr);
+        }
+      }
+      template <class parT_features>
+      void parT_tensor_filler(std::vector<float>& vdata,
+                              const parT::InputFeatures ifeature,
+                              const std::vector<parT_features>& 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<float>& inputs = inputs_parT(features.at(count), ifeature);
+          vdata.insert(vdata.end(), inputs.begin(), inputs.end());
+        }
+        unsigned int n_features = parT::N_InputFeatures.at(ifeature);
+        if (n < target_n)
+          vdata.insert(vdata.end(), (target_n - n) * n_features, 0);  // Add 0 to unfilled part as padding value
+      }
     }  // namespace btagbtvdeep
     #endif

RecoBTag/ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -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<edm::GlobalCache<ONNXRuntime>> {

    @@ -41,24 +44,9 @@ class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer
  
      std::vector<std::string> flav_names_;

      std::vector<std::string> input_names_;

      std::vector<std::string> 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<unsigned> input_sizes_;

      std::vector<std::vector<int64_t>> input_shapes_;  // shapes of each input group (-1 for dynamic axis)

    @@ -84,7 +72,7 @@ void ParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::Configurat
  
      desc.add<edm::InputTag>("src", edm::InputTag("pfParticleTransformerAK4TagInfos"));

      desc.add<std::vector<std::string>>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"});

      desc.add<edm::FileInPath>("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<std::vector<std::string>>("output_names", {"softmax"});

      desc.add<std::vector<std::string>>(

          "flav_names", std::vector<std::string>{"probb", "probbb", "problepb", "probc", "probuds", "probg"});

    @@ -124,12 +112,12 @@ void ParticleTransformerAK4ONNXJetTagsProducer::produce(edm::Event& iEvent, cons
  
          get_input_sizes(taginfo);

          // run prediction with dynamic batch size per event

          input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)n_features_cpf_},

                           {(int64_t)1, (int64_t)n_npf_, (int64_t)n_features_npf_},

                           {(int64_t)1, (int64_t)n_sv_, (int64_t)n_features_sv_},

                           {(int64_t)1, (int64_t)n_cpf_, (int64_t)n_pairwise_features_cpf_},

                           {(int64_t)1, (int64_t)n_npf_, (int64_t)n_pairwise_features_npf_},

                           {(int64_t)1, (int64_t)n_sv_, (int64_t)n_pairwise_features_sv_}};

          input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::N_InputFeatures.at(parT::kChargedCandidates)},

                           {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::N_InputFeatures.at(parT::kNeutralCandidates)},

                           {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::N_InputFeatures.at(parT::kVertices)},

                           {(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::N_InputFeatures.at(parT::kChargedCandidates4Vec)},

                           {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::N_InputFeatures.at(parT::kNeutralCandidates4Vec)},

                           {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::N_InputFeatures.at(parT::kVertices4Vec)}};

          outputs = globalCache()->run(input_names_, data_, input_shapes_, output_names_, 1)[0];

          assert(outputs.size() == flav_names_.size());

    @@ -151,24 +139,17 @@ void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
  
        const reco::FeaturesTagInfo<btagbtvdeep::ParticleTransformerAK4Features> taginfo) {

      const auto& features = taginfo.features();

      unsigned int n_cpf = features.c_pf_features.size();

      unsigned int n_npf = features.n_pf_features.size();

      unsigned int n_vtx = features.sv_features.size();

      n_cpf_ = std::clamp((unsigned int)features.c_pf_features.size(), (unsigned int)1, (unsigned int)parT::n_cpf_accept);

      n_npf_ = std::clamp((unsigned int)features.n_pf_features.size(), (unsigned int)1, (unsigned int)parT::n_npf_accept);

      n_sv_ = std::clamp((unsigned int)features.sv_features.size(), (unsigned int)1, (unsigned int)parT::n_sv_accept);

      n_cpf_ = std::max((unsigned int)1, n_cpf);

      n_npf_ = std::max((unsigned int)1, n_npf);

      n_sv_ = std::max((unsigned int)1, n_vtx);

      n_cpf_ = std::min((unsigned int)25, n_cpf_);

      n_npf_ = std::min((unsigned int)25, n_npf_);

      n_sv_ = std::min((unsigned int)5, n_sv_);

      input_sizes_ = {

          n_cpf_ * n_features_cpf_,

          n_npf_ * n_features_npf_,

          n_sv_ * n_features_sv_,

          n_cpf_ * n_pairwise_features_cpf_,

          n_npf_ * n_pairwise_features_npf_,

          n_sv_ * n_pairwise_features_sv_,

          n_cpf_ * parT::N_InputFeatures.at(parT::kChargedCandidates),

          n_npf_ * parT::N_InputFeatures.at(parT::kNeutralCandidates),

          n_sv_ * parT::N_InputFeatures.at(parT::kVertices),

          n_cpf_ * parT::N_InputFeatures.at(parT::kChargedCandidates4Vec),

          n_npf_ * parT::N_InputFeatures.at(parT::kNeutralCandidates4Vec),

          n_sv_ * parT::N_InputFeatures.at(parT::kVertices4Vec),

      };

      // init data storage

      data_.clear();

    @@ -180,116 +161,26 @@ void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
  
    }

    void ParticleTransformerAK4ONNXJetTagsProducer::make_inputs(btagbtvdeep::ParticleTransformerAK4Features features) {

      float* ptr = nullptr;

      //float* ptr = nullptr;

      const float* start = nullptr;

      unsigned offset = 0;

      // c_pf candidates

      auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);

      for (std::size_t c_pf_n = 0; c_pf_n < max_c_pf_n; c_pf_n++) {

        const auto& c_pf_features = features.c_pf_features.at(c_pf_n);

        ptr = &data_[kChargedCandidates][offset + c_pf_n * n_features_cpf_];

        start = ptr;

        *ptr = c_pf_features.btagPf_trackEtaRel;

        *(++ptr) = c_pf_features.btagPf_trackPtRel;

        *(++ptr) = c_pf_features.btagPf_trackPPar;

        *(++ptr) = c_pf_features.btagPf_trackDeltaR;

        *(++ptr) = c_pf_features.btagPf_trackPParRatio;

        *(++ptr) = c_pf_features.btagPf_trackSip2dVal;

        *(++ptr) = c_pf_features.btagPf_trackSip2dSig;

        *(++ptr) = c_pf_features.btagPf_trackSip3dVal;

        *(++ptr) = c_pf_features.btagPf_trackSip3dSig;

        *(++ptr) = c_pf_features.btagPf_trackJetDistVal;

        *(++ptr) = c_pf_features.ptrel;

        *(++ptr) = c_pf_features.drminsv;

        *(++ptr) = c_pf_features.vtx_ass;

        *(++ptr) = c_pf_features.puppiw;

        *(++ptr) = c_pf_features.chi2;

        *(++ptr) = c_pf_features.quality;

        assert(start + n_features_cpf_ - 1 == ptr);

      }

      // n_pf candidates

      auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);

      for (std::size_t n_pf_n = 0; n_pf_n < max_n_pf_n; n_pf_n++) {

        const auto& n_pf_features = features.n_pf_features.at(n_pf_n);

        ptr = &data_[kNeutralCandidates][offset + n_pf_n * n_features_npf_];

        start = ptr;

        *ptr = n_pf_features.ptrel;

        *(++ptr) = n_pf_features.etarel;

        *(++ptr) = n_pf_features.phirel;

        *(++ptr) = n_pf_features.deltaR;

        *(++ptr) = n_pf_features.isGamma;

        *(++ptr) = n_pf_features.hadFrac;

        *(++ptr) = n_pf_features.drminsv;

        *(++ptr) = n_pf_features.puppiw;

        assert(start + n_features_npf_ - 1 == ptr);

      }

      // sv candidates

      auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_);

      for (std::size_t sv_n = 0; sv_n < max_sv_n; sv_n++) {

        const auto& sv_features = features.sv_features.at(sv_n);

        ptr = &data_[kVertices][offset + sv_n * n_features_sv_];

        start = ptr;

        *ptr = sv_features.pt;

        *(++ptr) = sv_features.deltaR;

        *(++ptr) = sv_features.mass;

        *(++ptr) = sv_features.etarel;

        *(++ptr) = sv_features.phirel;

        *(++ptr) = sv_features.ntracks;

        *(++ptr) = sv_features.chi2;

        *(++ptr) = sv_features.normchi2;

        *(++ptr) = sv_features.dxy;

        *(++ptr) = sv_features.dxysig;

        *(++ptr) = sv_features.d3d;

        *(++ptr) = sv_features.d3dsig;

        *(++ptr) = sv_features.costhetasvpv;

        *(++ptr) = sv_features.enratio;

        assert(start + n_features_sv_ - 1 == ptr);

      }

      // c_pf candidates

      parT_tensor_filler(data_, parT::kChargedCandidates, features.c_pf_features, max_c_pf_n, start, offset);

      // n_pf candidates

      parT_tensor_filler(data_, parT::kNeutralCandidates, features.n_pf_features, max_n_pf_n, start, offset);

      // sv candidates

      parT_tensor_filler(data_, parT::kVertices, features.sv_features, max_sv_n, start, offset);

      // cpf pairwise features (4-vectors)

      auto max_cpf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);

      for (std::size_t cpf_n = 0; cpf_n < max_cpf_n; cpf_n++) {

        const auto& cpf_pairwise_features = features.c_pf_features.at(cpf_n);

        ptr = &data_[kChargedCandidates4Vec][offset + cpf_n * n_pairwise_features_cpf_];

        start = ptr;

        *ptr = cpf_pairwise_features.px;

        *(++ptr) = cpf_pairwise_features.py;

        *(++ptr) = cpf_pairwise_features.pz;

        *(++ptr) = cpf_pairwise_features.e;

        assert(start + n_pairwise_features_cpf_ - 1 == ptr);

      }

      parT_tensor_filler(data_, parT::kChargedCandidates4Vec, features.c_pf_features, max_c_pf_n, start, offset);

      // npf pairwise features (4-vectors)

      auto max_npf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);

      for (std::size_t npf_n = 0; npf_n < max_npf_n; npf_n++) {

        const auto& npf_pairwise_features = features.n_pf_features.at(npf_n);

        ptr = &data_[kNeutralCandidates4Vec][offset + npf_n * n_pairwise_features_npf_];

        start = ptr;

        *ptr = npf_pairwise_features.px;

        *(++ptr) = npf_pairwise_features.py;

        *(++ptr) = npf_pairwise_features.pz;

        *(++ptr) = npf_pairwise_features.e;

        assert(start + n_pairwise_features_npf_ - 1 == ptr);

      }

      parT_tensor_filler(data_, parT::kNeutralCandidates4Vec, features.n_pf_features, max_n_pf_n, start, offset);

      // sv pairwise features (4-vectors)

      auto max_sv_N = std::min(features.sv_features.size(), (std::size_t)n_sv_);

      for (std::size_t sv_N = 0; sv_N < max_sv_N; sv_N++) {

        const auto& sv_pairwise_features = features.sv_features.at(sv_N);

        ptr = &data_[kVertices4Vec][offset + sv_N * n_pairwise_features_sv_];

        start = ptr;

        *ptr = sv_pairwise_features.px;

        *(++ptr) = sv_pairwise_features.py;

        *(++ptr) = sv_pairwise_features.pz;

        *(++ptr) = sv_pairwise_features.e;

        assert(start + n_pairwise_features_sv_ - 1 == ptr);

      }

      parT_tensor_filler(data_, parT::kVertices4Vec, features.sv_features, max_sv_n, start, offset);

    }

    //define this as a plug-in

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Particle Transformer sonic producer #45579

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!