ONNXRuntime/plugins/ParticleTransformerAK4ONNXJetTagsProducer.cc

0001 #include "FWCore/Framework/interface/Frameworkfwd.h"
0002 #include "FWCore/Framework/interface/stream/EDProducer.h"
0003
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "FWCore/Framework/interface/MakerMacros.h"
0006
0007 #include "FWCore/Framework/interface/makeRefToBaseProdFrom.h"
0008
0009 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0010 #include "FWCore/Utilities/interface/StreamID.h"
0011
0012 #include "DataFormats/BTauReco/interface/JetTag.h"
0013
0014 #include "DataFormats/BTauReco/interface/ParticleTransformerAK4TagInfo.h"
0015 #include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h"
0016
0017 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0018
0019 using namespace cms::Ort;
0020
0021 class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer<edm::GlobalCache<ONNXRuntime>> {
0022 public:
0023   explicit ParticleTransformerAK4ONNXJetTagsProducer(const edm::ParameterSet&, const ONNXRuntime*);
0024   ~ParticleTransformerAK4ONNXJetTagsProducer() override = default;
0025
0026   static void fillDescriptions(edm::ConfigurationDescriptions&);
0027
0028   static std::unique_ptr<ONNXRuntime> initializeGlobalCache(const edm::ParameterSet&);
0029   static void globalEndJob(const ONNXRuntime*);
0030
0031 private:
0032   typedef std::vector<reco::ParticleTransformerAK4TagInfo> TagInfoCollection;
0033   typedef reco::JetTagCollection JetTagCollection;
0034
0035   void produce(edm::Event&, const edm::EventSetup&) override;
0036
0037   void make_inputs(btagbtvdeep::ParticleTransformerAK4Features features);
0038   void get_input_sizes(const reco::FeaturesTagInfo<btagbtvdeep::ParticleTransformerAK4Features> taginfo);
0039
0040   const edm::EDGetTokenT<TagInfoCollection> src_;
0041   std::vector<std::string> flav_names_;
0042   std::vector<std::string> input_names_;
0043   std::vector<std::string> output_names_;
0044
0045   enum InputIndexes {
0046     kChargedCandidates = 0,
0047     kNeutralCandidates = 1,
0048     kVertices = 2,
0049     kChargedCandidates4Vec = 3,
0050     kNeutralCandidates4Vec = 4,
0051     kVertices4Vec = 5
0052   };
0053   unsigned n_cpf_;
0054   constexpr static unsigned n_features_cpf_ = 16;
0055   constexpr static unsigned n_pairwise_features_cpf_ = 4;
0056   unsigned n_npf_;
0057   constexpr static unsigned n_features_npf_ = 8;
0058   constexpr static unsigned n_pairwise_features_npf_ = 4;
0059   unsigned n_sv_;
0060   constexpr static unsigned n_features_sv_ = 14;
0061   constexpr static unsigned n_pairwise_features_sv_ = 4;
0062   std::vector<unsigned> input_sizes_;
0063   std::vector<std::vector<int64_t>> input_shapes_;  // shapes of each input group (-1 for dynamic axis)
0064
0065   // hold the input data
0066   FloatArrays data_;
0067 };
0068
0069 ParticleTransformerAK4ONNXJetTagsProducer::ParticleTransformerAK4ONNXJetTagsProducer(const edm::ParameterSet& iConfig,
0070                                                                                      const ONNXRuntime* cache)
0071     : src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
0072       flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
0073       input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
0074       output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")) {
0075   // get output names from flav_names
0076   for (const auto& flav_name : flav_names_) {
0077     produces<JetTagCollection>(flav_name);
0078   }
0079 }
0080
0081 void ParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0082   // pfParticleTransformerAK4JetTags
0083   edm::ParameterSetDescription desc;
0084   desc.add<edm::InputTag>("src", edm::InputTag("pfParticleTransformerAK4TagInfos"));
0085   desc.add<std::vector<std::string>>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"});
0086   desc.add<edm::FileInPath>("model_path",
0087                             edm::FileInPath("RecoBTag/Combined/data/RobustParTAK4/PUPPI/V00/RobustParTAK4.onnx"));
0088   desc.add<std::vector<std::string>>("output_names", {"softmax"});
0089   desc.add<std::vector<std::string>>(
0090       "flav_names", std::vector<std::string>{"probb", "probbb", "problepb", "probc", "probuds", "probg"});
0091
0092   descriptions.add("pfParticleTransformerAK4JetTags", desc);
0093 }
0094
0095 std::unique_ptr<ONNXRuntime> ParticleTransformerAK4ONNXJetTagsProducer::initializeGlobalCache(
0096     const edm::ParameterSet& iConfig) {
0097   return std::make_unique<ONNXRuntime>(iConfig.getParameter<edm::FileInPath>("model_path").fullPath());
0098 }
0099
0100 void ParticleTransformerAK4ONNXJetTagsProducer::globalEndJob(const ONNXRuntime* cache) {}
0101
0102 void ParticleTransformerAK4ONNXJetTagsProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0103   edm::Handle<TagInfoCollection> tag_infos;
0104   iEvent.getByToken(src_, tag_infos);
0105
0106   // initialize output collection
0107   std::vector<std::unique_ptr<JetTagCollection>> output_tags;
0108   if (!tag_infos->empty()) {
0109     auto jet_ref = tag_infos->begin()->jet();
0110     auto ref2prod = edm::makeRefToBaseProdFrom(jet_ref, iEvent);
0111     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0112       output_tags.emplace_back(std::make_unique<JetTagCollection>(ref2prod));
0113     }
0114   } else {
0115     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0116       output_tags.emplace_back(std::make_unique<JetTagCollection>());
0117     }
0118   }
0119
0120   for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0121     const auto& taginfo = (*tag_infos)[jet_n];
0122     std::vector<float> outputs(flav_names_.size(), -1.0);
0123     if (taginfo.features().is_filled) {
0124       get_input_sizes(taginfo);
0125
0126       // run prediction with dynamic batch size per event
0127       input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)n_features_cpf_},
0128                        {(int64_t)1, (int64_t)n_npf_, (int64_t)n_features_npf_},
0129                        {(int64_t)1, (int64_t)n_sv_, (int64_t)n_features_sv_},
0130                        {(int64_t)1, (int64_t)n_cpf_, (int64_t)n_pairwise_features_cpf_},
0131                        {(int64_t)1, (int64_t)n_npf_, (int64_t)n_pairwise_features_npf_},
0132                        {(int64_t)1, (int64_t)n_sv_, (int64_t)n_pairwise_features_sv_}};
0133
0134       outputs = globalCache()->run(input_names_, data_, input_shapes_, output_names_, 1)[0];
0135       assert(outputs.size() == flav_names_.size());
0136     }
0137
0138     const auto& jet_ref = tag_infos->at(jet_n).jet();
0139     for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0140       (*(output_tags[flav_n]))[jet_ref] = outputs[flav_n];
0141     }
0142   }
0143
0144   // put into the event
0145   for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) {
0146     iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]);
0147   }
0148 }
0149
0150 void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
0151     const reco::FeaturesTagInfo<btagbtvdeep::ParticleTransformerAK4Features> taginfo) {
0152   const auto& features = taginfo.features();
0153
0154   unsigned int n_cpf = features.c_pf_features.size();
0155   unsigned int n_npf = features.n_pf_features.size();
0156   unsigned int n_vtx = features.sv_features.size();
0157
0158   n_cpf_ = std::max((unsigned int)1, n_cpf);
0159   n_npf_ = std::max((unsigned int)1, n_npf);
0160   n_sv_ = std::max((unsigned int)1, n_vtx);
0161
0162   n_cpf_ = std::min((unsigned int)25, n_cpf_);
0163   n_npf_ = std::min((unsigned int)25, n_npf_);
0164   n_sv_ = std::min((unsigned int)5, n_sv_);
0165   input_sizes_ = {
0166       n_cpf_ * n_features_cpf_,
0167       n_npf_ * n_features_npf_,
0168       n_sv_ * n_features_sv_,
0169       n_cpf_ * n_pairwise_features_cpf_,
0170       n_npf_ * n_pairwise_features_npf_,
0171       n_sv_ * n_pairwise_features_sv_,
0172   };
0173   // init data storage
0174   data_.clear();
0175   for (const auto& len : input_sizes_) {
0176     data_.emplace_back(1 * len, 0);
0177   }
0178
0179   make_inputs(features);
0180 }
0181
0182 void ParticleTransformerAK4ONNXJetTagsProducer::make_inputs(btagbtvdeep::ParticleTransformerAK4Features features) {
0183   float* ptr = nullptr;
0184   const float* start = nullptr;
0185   unsigned offset = 0;
0186
0187   // c_pf candidates
0188   auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
0189   for (std::size_t c_pf_n = 0; c_pf_n < max_c_pf_n; c_pf_n++) {
0190     const auto& c_pf_features = features.c_pf_features.at(c_pf_n);
0191     ptr = &data_[kChargedCandidates][offset + c_pf_n * n_features_cpf_];
0192     start = ptr;
0193     *ptr = c_pf_features.btagPf_trackEtaRel;
0194     *(++ptr) = c_pf_features.btagPf_trackPtRel;
0195     *(++ptr) = c_pf_features.btagPf_trackPPar;
0196     *(++ptr) = c_pf_features.btagPf_trackDeltaR;
0197     *(++ptr) = c_pf_features.btagPf_trackPParRatio;
0198     *(++ptr) = c_pf_features.btagPf_trackSip2dVal;
0199     *(++ptr) = c_pf_features.btagPf_trackSip2dSig;
0200     *(++ptr) = c_pf_features.btagPf_trackSip3dVal;
0201     *(++ptr) = c_pf_features.btagPf_trackSip3dSig;
0202     *(++ptr) = c_pf_features.btagPf_trackJetDistVal;
0203     *(++ptr) = c_pf_features.ptrel;
0204     *(++ptr) = c_pf_features.drminsv;
0205     *(++ptr) = c_pf_features.vtx_ass;
0206     *(++ptr) = c_pf_features.puppiw;
0207     *(++ptr) = c_pf_features.chi2;
0208     *(++ptr) = c_pf_features.quality;
0209     assert(start + n_features_cpf_ - 1 == ptr);
0210   }
0211
0212   // n_pf candidates
0213   auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
0214   for (std::size_t n_pf_n = 0; n_pf_n < max_n_pf_n; n_pf_n++) {
0215     const auto& n_pf_features = features.n_pf_features.at(n_pf_n);
0216     ptr = &data_[kNeutralCandidates][offset + n_pf_n * n_features_npf_];
0217     start = ptr;
0218     *ptr = n_pf_features.ptrel;
0219     *(++ptr) = n_pf_features.etarel;
0220     *(++ptr) = n_pf_features.phirel;
0221     *(++ptr) = n_pf_features.deltaR;
0222     *(++ptr) = n_pf_features.isGamma;
0223     *(++ptr) = n_pf_features.hadFrac;
0224     *(++ptr) = n_pf_features.drminsv;
0225     *(++ptr) = n_pf_features.puppiw;
0226     assert(start + n_features_npf_ - 1 == ptr);
0227   }
0228
0229   // sv candidates
0230   auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_);
0231   for (std::size_t sv_n = 0; sv_n < max_sv_n; sv_n++) {
0232     const auto& sv_features = features.sv_features.at(sv_n);
0233     ptr = &data_[kVertices][offset + sv_n * n_features_sv_];
0234     start = ptr;
0235     *ptr = sv_features.pt;
0236     *(++ptr) = sv_features.deltaR;
0237     *(++ptr) = sv_features.mass;
0238     *(++ptr) = sv_features.etarel;
0239     *(++ptr) = sv_features.phirel;
0240     *(++ptr) = sv_features.ntracks;
0241     *(++ptr) = sv_features.chi2;
0242     *(++ptr) = sv_features.normchi2;
0243     *(++ptr) = sv_features.dxy;
0244     *(++ptr) = sv_features.dxysig;
0245     *(++ptr) = sv_features.d3d;
0246     *(++ptr) = sv_features.d3dsig;
0247     *(++ptr) = sv_features.costhetasvpv;
0248     *(++ptr) = sv_features.enratio;
0249     assert(start + n_features_sv_ - 1 == ptr);
0250   }
0251
0252   // cpf pairwise features (4-vectors)
0253   auto max_cpf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
0254   for (std::size_t cpf_n = 0; cpf_n < max_cpf_n; cpf_n++) {
0255     const auto& cpf_pairwise_features = features.c_pf_features.at(cpf_n);
0256     ptr = &data_[kChargedCandidates4Vec][offset + cpf_n * n_pairwise_features_cpf_];
0257     start = ptr;
0258     *ptr = cpf_pairwise_features.px;
0259     *(++ptr) = cpf_pairwise_features.py;
0260     *(++ptr) = cpf_pairwise_features.pz;
0261     *(++ptr) = cpf_pairwise_features.e;
0262
0263     assert(start + n_pairwise_features_cpf_ - 1 == ptr);
0264   }
0265
0266   // npf pairwise features (4-vectors)
0267   auto max_npf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
0268   for (std::size_t npf_n = 0; npf_n < max_npf_n; npf_n++) {
0269     const auto& npf_pairwise_features = features.n_pf_features.at(npf_n);
0270     ptr = &data_[kNeutralCandidates4Vec][offset + npf_n * n_pairwise_features_npf_];
0271     start = ptr;
0272     *ptr = npf_pairwise_features.px;
0273     *(++ptr) = npf_pairwise_features.py;
0274     *(++ptr) = npf_pairwise_features.pz;
0275     *(++ptr) = npf_pairwise_features.e;
0276
0277     assert(start + n_pairwise_features_npf_ - 1 == ptr);
0278   }
0279
0280   // sv pairwise features (4-vectors)
0281   auto max_sv_N = std::min(features.sv_features.size(), (std::size_t)n_sv_);
0282   for (std::size_t sv_N = 0; sv_N < max_sv_N; sv_N++) {
0283     const auto& sv_pairwise_features = features.sv_features.at(sv_N);
0284     ptr = &data_[kVertices4Vec][offset + sv_N * n_pairwise_features_sv_];
0285     start = ptr;
0286     *ptr = sv_pairwise_features.px;
0287     *(++ptr) = sv_pairwise_features.py;
0288     *(++ptr) = sv_pairwise_features.pz;
0289     *(++ptr) = sv_pairwise_features.e;
0290
0291     assert(start + n_pairwise_features_sv_ - 1 == ptr);
0292   }
0293 }
0294
0295 //define this as a plug-in
0296 DEFINE_FWK_MODULE(ParticleTransformerAK4ONNXJetTagsProducer);