ONNXRuntime/plugins/ParticleTransformerAK4SonicJetTagsProducer.cc

0001 #include "FWCore/Framework/interface/Frameworkfwd.h"
0002
0003 #include "FWCore/Framework/interface/Event.h"
0004 #include "FWCore/Framework/interface/MakerMacros.h"
0005
0006 #include "FWCore/Framework/interface/makeRefToBaseProdFrom.h"
0007
0008 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0009 #include "FWCore/Utilities/interface/StreamID.h"
0010
0011 #include "DataFormats/BTauReco/interface/JetTag.h"
0012
0013 #include "DataFormats/BTauReco/interface/ParticleTransformerAK4TagInfo.h"
0014 #include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h"
0015
0016 #include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h"
0017 #include "HeterogeneousCore/SonicTriton/interface/TritonData.h"
0018
0019 #include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h"
0020 #include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
0021
0022 class ParticleTransformerAK4SonicJetTagsProducer : public TritonEDProducer<> {
0023 public:
0024   explicit ParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet &);
0025   ~ParticleTransformerAK4SonicJetTagsProducer() override;
0026
0027   void acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) override;
0028
0029   void produce(edm::Event &iEvent, edm::EventSetup const &iSetup, Output const &iOutput) override;
0030   static void fillDescriptions(edm::ConfigurationDescriptions &);
0031
0032 private:
0033   typedef std::vector<reco::ParticleTransformerAK4TagInfo> TagInfoCollection;
0034   typedef reco::JetTagCollection JetTagCollection;
0035
0036   const edm::EDGetTokenT<TagInfoCollection> src_;
0037   std::vector<std::string> flav_names_;
0038   std::vector<std::string> input_names_;
0039   std::vector<std::string> output_names_;
0040
0041   bool skippedInference_ = false;
0042 };
0043
0044 ParticleTransformerAK4SonicJetTagsProducer::ParticleTransformerAK4SonicJetTagsProducer(const edm::ParameterSet &iConfig)
0045     : TritonEDProducer<>(iConfig),
0046       src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
0047       flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
0048       input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
0049       output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")) {
0050   // get output names from flav_names
0051   for (const auto &flav_name : flav_names_) {
0052     produces<JetTagCollection>(flav_name);
0053   }
0054 }
0055
0056 ParticleTransformerAK4SonicJetTagsProducer::~ParticleTransformerAK4SonicJetTagsProducer() {}
0057
0058 void ParticleTransformerAK4SonicJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0059   // pfParticleTransformerAK4JetTags
0060   edm::ParameterSetDescription desc;
0061   TritonClient::fillPSetDescription(desc);
0062   desc.add<edm::InputTag>("src", edm::InputTag("pfParticleTransformerAK4TagInfos"));
0063   desc.add<std::vector<std::string>>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"});
0064   desc.add<std::vector<std::string>>("output_names", {"softmax"});
0065   desc.add<std::vector<std::string>>(
0066       "flav_names", std::vector<std::string>{"probb", "probbb", "problepb", "probc", "probuds", "probg"});
0067
0068   descriptions.add("pfParticleTransformerAK4SonicJetTags", desc);
0069 }
0070
0071 void ParticleTransformerAK4SonicJetTagsProducer::acquire(edm::Event const &iEvent,
0072                                                          edm::EventSetup const &iSetup,
0073                                                          Input &iInput) {
0074   edm::Handle<TagInfoCollection> tag_infos;
0075   iEvent.getByToken(src_, tag_infos);
0076   client_->setBatchSize(tag_infos->size());
0077   skippedInference_ = false;
0078   if (tag_infos->empty())
0079     return;
0080
0081   // Find the max n_cpf, n_npf and n_vtx among all the jets in an event.
0082   unsigned int max_n_cpf_counter = 0;
0083   unsigned int max_n_npf_counter = 0;
0084   unsigned int max_n_vtx_counter = 0;
0085   for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0086     max_n_cpf_counter =
0087         std::max(max_n_cpf_counter, static_cast<unsigned int>(((*tag_infos)[jet_n]).features().c_pf_features.size()));
0088     max_n_npf_counter =
0089         std::max(max_n_npf_counter, static_cast<unsigned int>(((*tag_infos)[jet_n]).features().n_pf_features.size()));
0090     max_n_vtx_counter =
0091         std::max(max_n_vtx_counter, static_cast<unsigned int>(((*tag_infos)[jet_n]).features().sv_features.size()));
0092   }
0093
0094   // If an event has no jet, or all jets has zero n_cpf, n_npf and n_vtx, the inference is skipped.
0095   if (max_n_cpf_counter == 0 && max_n_npf_counter == 0 && max_n_vtx_counter == 0) {
0096     client_->setBatchSize(0);
0097     skippedInference_ = true;
0098     return;
0099   }
0100
0101   // all the jets in the same event will fill up the same amount of n_cpf, n_npf, n_vtx and send to server
0102   const unsigned int target_n_cpf = std::clamp(max_n_cpf_counter, (unsigned int)1, (unsigned int)parT::n_cpf_accept);
0103   const unsigned int target_n_npf = std::clamp(max_n_npf_counter, (unsigned int)1, (unsigned int)parT::n_npf_accept);
0104   const unsigned int target_n_vtx = std::clamp(max_n_vtx_counter, (unsigned int)1, (unsigned int)parT::n_sv_accept);
0105
0106   const std::map<parT::InputFeatures, unsigned int> target_n{{parT::kChargedCandidates, target_n_cpf},
0107                                                              {parT::kNeutralCandidates, target_n_npf},
0108                                                              {parT::kVertices, target_n_vtx},
0109                                                              {parT::kChargedCandidates4Vec, target_n_cpf},
0110                                                              {parT::kNeutralCandidates4Vec, target_n_npf},
0111                                                              {parT::kVertices4Vec, target_n_vtx}};
0112
0113   // loop through all groups of features
0114   for (parT::InputFeatures ifeature = parT::kBegin; ifeature != parT::kEnd;
0115        ifeature = static_cast<parT::InputFeatures>(ifeature + 1)) {
0116     const auto &group_name = input_names_[ifeature];
0117     auto &input = iInput.at(group_name);
0118
0119     input.setShape(0, target_n.at(ifeature));
0120     auto tdata = input.allocate<float>(true);
0121
0122     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0123       const auto &taginfo = (*tag_infos)[jet_n];
0124       const auto &features = taginfo.features();
0125       auto &vdata = (*tdata)[jet_n];
0126       if (ifeature == parT::kChargedCandidates || ifeature == parT::kChargedCandidates4Vec)
0127         parT_tensor_filler(vdata, ifeature, features.c_pf_features, target_n_cpf);
0128       else if (ifeature == parT::kNeutralCandidates || ifeature == parT::kNeutralCandidates4Vec)
0129         parT_tensor_filler(vdata, ifeature, features.n_pf_features, target_n_npf);
0130       else if (ifeature == parT::kVertices || ifeature == parT::kVertices4Vec)
0131         parT_tensor_filler(vdata, ifeature, features.sv_features, target_n_vtx);
0132     }
0133     input.toServer(tdata);
0134   }
0135 }
0136
0137 void ParticleTransformerAK4SonicJetTagsProducer::produce(edm::Event &iEvent,
0138                                                          const edm::EventSetup &iSetup,
0139                                                          Output const &iOutput) {
0140   edm::Handle<TagInfoCollection> tag_infos;
0141   iEvent.getByToken(src_, tag_infos);
0142
0143   // initialize output collection
0144   std::vector<std::unique_ptr<JetTagCollection>> output_tags;
0145   if (!tag_infos->empty()) {
0146     auto jet_ref = tag_infos->begin()->jet();
0147     auto ref2prod = edm::makeRefToBaseProdFrom(jet_ref, iEvent);
0148     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0149       output_tags.emplace_back(std::make_unique<JetTagCollection>(ref2prod));
0150     }
0151   } else {
0152     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0153       output_tags.emplace_back(std::make_unique<JetTagCollection>());
0154     }
0155   }
0156   if (!tag_infos->empty()) {
0157     if (!skippedInference_) {
0158       const auto &output1 = iOutput.begin()->second;
0159       const auto &outputs_from_server = output1.fromServer<float>();
0160
0161       for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0162         const auto &taginfo = (*tag_infos)[jet_n];
0163         const auto &jet_ref = tag_infos->at(jet_n).jet();
0164
0165         if (taginfo.features().is_filled) {
0166           for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0167             (*(output_tags[flav_n]))[jet_ref] = outputs_from_server[jet_n][flav_n];
0168           }
0169         } else {
0170           for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0171             (*(output_tags[flav_n]))[jet_ref] = -1.0;
0172           }
0173         }
0174       }
0175     } else {
0176       for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0177         const auto &jet_ref = tag_infos->at(jet_n).jet();
0178         for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0179           (*(output_tags[flav_n]))[jet_ref] = -1.0;
0180         }
0181       }
0182     }
0183   }
0184   // put into the event
0185   for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) {
0186     iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]);
0187   }
0188 }
0189
0190 //define this as a plug-in
0191 DEFINE_FWK_MODULE(ParticleTransformerAK4SonicJetTagsProducer);