ONNXRuntime/plugins/DeepCombinedONNXJetTagsProducer.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/DeepFlavourTagInfo.h"
0015
0016 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0017
0018 #include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h"
0019 #include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
0020
0021 using namespace cms::Ort;
0022
0023 class DeepCombinedONNXJetTagsProducer : public edm::stream::EDProducer<edm::GlobalCache<ONNXRuntime>> {
0024 public:
0025   explicit DeepCombinedONNXJetTagsProducer(const edm::ParameterSet&, const ONNXRuntime*);
0026   ~DeepCombinedONNXJetTagsProducer() override;
0027
0028   static void fillDescriptions(edm::ConfigurationDescriptions&);
0029
0030   static std::unique_ptr<ONNXRuntime> initializeGlobalCache(const edm::ParameterSet&);
0031   static void globalEndJob(const ONNXRuntime*);
0032
0033 private:
0034   typedef std::vector<reco::DeepFlavourTagInfo> TagInfoCollection;
0035   typedef reco::JetTagCollection JetTagCollection;
0036
0037   void produce(edm::Event&, const edm::EventSetup&) override;
0038
0039   void make_inputs(unsigned i_jet, const reco::DeepFlavourTagInfo& taginfo);
0040
0041   const edm::EDGetTokenT<TagInfoCollection> src_;
0042   std::vector<std::string> flav_names_;
0043   std::vector<std::string> input_names_;
0044   std::vector<std::string> output_names_;
0045
0046   const double min_jet_pt_;
0047   const double max_jet_eta_;
0048
0049   enum InputIndexes {
0050     kGlobal = 0,
0051     kChargedCandidates = 1,
0052     kNeutralCandidates = 2,
0053     kVertices = 3,
0054     kGlobal1 = 4,
0055     kSeedingTracks = 5,
0056     kNeighbourTracks = 6
0057   };
0058   const static unsigned n_features_global_ = deepflavour::n_features_global;
0059   const static unsigned n_cpf_ = deepflavour::n_cpf;
0060   const static unsigned n_features_cpf_ = deepflavour::n_features_cpf;
0061   const static unsigned n_npf_ = deepflavour::n_npf;
0062   const static unsigned n_features_npf_ = deepflavour::n_features_npf;
0063   const static unsigned n_sv_ = deepflavour::n_sv;
0064   const static unsigned n_features_sv_ = deepflavour::n_features_sv;
0065   const static unsigned n_features_global1_ = deepvertex::n_features_global;
0066   const static unsigned n_seed_ = deepvertex::n_seed;
0067   const static unsigned n_features_seed_ = deepvertex::n_features_seed;
0068   const static unsigned n_neighbor_ = deepvertex::n_neighbor;
0069   const static unsigned n_features_neighbor_ = deepvertex::n_features_neighbor;
0070
0071   const static std::vector<unsigned> input_sizes_;
0072
0073   // hold the input data
0074   FloatArrays data_;
0075 };
0076
0077 const std::vector<unsigned> DeepCombinedONNXJetTagsProducer::input_sizes_{n_features_global_,
0078                                                                           n_cpf_* n_features_cpf_,
0079                                                                           n_npf_* n_features_npf_,
0080                                                                           n_sv_* n_features_sv_,
0081                                                                           n_features_global1_,
0082                                                                           n_seed_* n_features_seed_,
0083                                                                           n_neighbor_* n_features_neighbor_,
0084                                                                           n_neighbor_* n_features_neighbor_,
0085                                                                           n_neighbor_* n_features_neighbor_,
0086                                                                           n_neighbor_* n_features_neighbor_,
0087                                                                           n_neighbor_* n_features_neighbor_,
0088                                                                           n_neighbor_* n_features_neighbor_,
0089                                                                           n_neighbor_* n_features_neighbor_,
0090                                                                           n_neighbor_* n_features_neighbor_,
0091                                                                           n_neighbor_* n_features_neighbor_,
0092                                                                           n_neighbor_* n_features_neighbor_};
0093
0094 DeepCombinedONNXJetTagsProducer::DeepCombinedONNXJetTagsProducer(const edm::ParameterSet& iConfig,
0095                                                                  const ONNXRuntime* cache)
0096     : src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
0097       flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
0098       input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
0099       output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")),
0100       min_jet_pt_(iConfig.getParameter<double>("min_jet_pt")),
0101       max_jet_eta_(iConfig.getParameter<double>("max_jet_eta")) {
0102   // get output names from flav_names
0103   for (const auto& flav_name : flav_names_) {
0104     produces<JetTagCollection>(flav_name);
0105   }
0106
0107   assert(input_names_.size() == input_sizes_.size());
0108 }
0109
0110 DeepCombinedONNXJetTagsProducer::~DeepCombinedONNXJetTagsProducer() {}
0111
0112 void DeepCombinedONNXJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0113   // pfDeepFlavourJetTags
0114   edm::ParameterSetDescription desc;
0115   desc.add<edm::InputTag>("src", edm::InputTag("pfDeepFlavourTagInfos"));
0116   desc.add<std::vector<std::string>>("input_names",
0117                                      {"input_1_DFla",
0118                                       "input_2_DFla",
0119                                       "input_3_DFla",
0120                                       "input_4_DFla",
0121                                       "input_1",
0122                                       "input_2",
0123                                       "input_3",
0124                                       "input_4",
0125                                       "input_5",
0126                                       "input_6",
0127                                       "input_7",
0128                                       "input_8",
0129                                       "input_9",
0130                                       "input_10",
0131                                       "input_11",
0132                                       "input_12"});
0133   desc.add<edm::FileInPath>("model_path",
0134                             edm::FileInPath("RecoBTag/Combined/data/DeepVertex/phase1_deepvertexcombined.onnx"));
0135   desc.add<std::vector<std::string>>("output_names", {"dense_13"});
0136   desc.add<std::vector<std::string>>("flav_names", std::vector<std::string>{"probb", "probc", "probuds", "probg"});
0137   desc.add<double>("min_jet_pt", 15.0);
0138   desc.add<double>("max_jet_eta", 2.5);
0139
0140   descriptions.add("pfDeepCombinedJetTags", desc);
0141 }
0142
0143 std::unique_ptr<ONNXRuntime> DeepCombinedONNXJetTagsProducer::initializeGlobalCache(const edm::ParameterSet& iConfig) {
0144   return std::make_unique<ONNXRuntime>(iConfig.getParameter<edm::FileInPath>("model_path").fullPath());
0145 }
0146
0147 void DeepCombinedONNXJetTagsProducer::globalEndJob(const ONNXRuntime* cache) {}
0148
0149 void DeepCombinedONNXJetTagsProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0150   edm::Handle<TagInfoCollection> tag_infos;
0151   iEvent.getByToken(src_, tag_infos);
0152
0153   data_.clear();
0154
0155   std::vector<std::unique_ptr<JetTagCollection>> output_tags;
0156   if (!tag_infos->empty()) {
0157     unsigned good_taginfo_count = 0;
0158     std::vector<bool> good_taginfo_jets(tag_infos->size(), false);
0159     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0160       const auto& jet_ref = (*tag_infos)[jet_n].jet();
0161       if (jet_ref->pt() > min_jet_pt_ && std::fabs(jet_ref->eta()) < max_jet_eta_) {
0162         good_taginfo_count++;
0163         good_taginfo_jets[jet_n] = true;
0164       }
0165     }
0166
0167     // init data storage w correct size
0168     for (const auto& len : input_sizes_) {
0169       data_.emplace_back(good_taginfo_count * len, 0);
0170     }
0171
0172     // initialize output collection
0173     auto jet_ref = tag_infos->begin()->jet();
0174     auto ref2prod = edm::makeRefToBaseProdFrom(jet_ref, iEvent);
0175     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0176       output_tags.emplace_back(std::make_unique<JetTagCollection>(ref2prod));
0177     }
0178
0179     // convert inputs
0180     unsigned inputs_done_count = 0;
0181     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0182       if (good_taginfo_jets[jet_n]) {
0183         const auto& taginfo = (*tag_infos)[jet_n];
0184         make_inputs(inputs_done_count, taginfo);
0185         inputs_done_count++;
0186       }
0187     }
0188
0189     // run prediction
0190     assert(inputs_done_count == good_taginfo_count);
0191     const auto outputs = globalCache()->run(input_names_, data_, {}, output_names_, good_taginfo_count)[0];
0192     assert(outputs.size() == flav_names_.size() * good_taginfo_count);
0193
0194     // get the outputs
0195     unsigned i_output = 0;
0196     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0197       const auto& jet_ref = (*tag_infos)[jet_n].jet();
0198       for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0199         if (good_taginfo_jets[jet_n]) {
0200           (*(output_tags[flav_n]))[jet_ref] = outputs[i_output];
0201           ++i_output;
0202         } else {
0203           (*(output_tags[flav_n]))[jet_ref] = -2;
0204         }
0205       }
0206     }
0207   } else {
0208     // create empty output collection
0209     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0210       output_tags.emplace_back(std::make_unique<JetTagCollection>());
0211     }
0212   }
0213
0214   // put into the event
0215   for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) {
0216     iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]);
0217   }
0218 }
0219
0220 void DeepCombinedONNXJetTagsProducer::make_inputs(unsigned i_jet, const reco::DeepFlavourTagInfo& taginfo) {
0221   const auto& features = taginfo.features();
0222   float* ptr = nullptr;
0223   const float* start = nullptr;
0224   unsigned offset = 0;
0225
0226   // jet and other global features
0227   offset = i_jet * input_sizes_[kGlobal];
0228   ptr = &data_[kGlobal][offset];
0229   start = ptr;
0230   jet_tensor_filler(ptr, features);
0231   assert(start + n_features_global_ - 1 == ptr);
0232
0233   // c_pf candidates
0234   auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
0235   offset = i_jet * input_sizes_[kChargedCandidates];
0236   for (std::size_t c_pf_n = 0; c_pf_n < max_c_pf_n; c_pf_n++) {
0237     const auto& c_pf_features = features.c_pf_features[c_pf_n];
0238     ptr = &data_[kChargedCandidates][offset + c_pf_n * n_features_cpf_];
0239     start = ptr;
0240     cpf_tensor_filler(ptr, c_pf_features);
0241     assert(start + n_features_cpf_ - 1 == ptr);
0242   }
0243
0244   // n_pf candidates
0245   auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
0246   offset = i_jet * input_sizes_[kNeutralCandidates];
0247   for (std::size_t n_pf_n = 0; n_pf_n < max_n_pf_n; n_pf_n++) {
0248     const auto& n_pf_features = features.n_pf_features[n_pf_n];
0249     ptr = &data_[kNeutralCandidates][offset + n_pf_n * n_features_npf_];
0250     start = ptr;
0251     npf_tensor_filler(ptr, n_pf_features);
0252     assert(start + n_features_npf_ - 1 == ptr);
0253   }
0254
0255   // sv candidates
0256   auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_);
0257   offset = i_jet * input_sizes_[kVertices];
0258   for (std::size_t sv_n = 0; sv_n < max_sv_n; sv_n++) {
0259     const auto& sv_features = features.sv_features[sv_n];
0260     ptr = &data_[kVertices][offset + sv_n * n_features_sv_];
0261     start = ptr;
0262     sv_tensor_filler(ptr, sv_features);
0263     assert(start + n_features_sv_ - 1 == ptr);
0264   }
0265
0266   // jet variables
0267   offset = i_jet * input_sizes_[kGlobal1];
0268   const auto& jet_features = features.jet_features;
0269   ptr = &data_[kGlobal1][offset];
0270   start = ptr;
0271   jet4vec_tensor_filler(ptr, jet_features);
0272   assert(start + n_features_global1_ - 1 == ptr);
0273
0274   // seeds
0275   auto max_seed_n = std::min(features.seed_features.size(), (std::size_t)n_seed_);
0276   offset = i_jet * input_sizes_[kSeedingTracks];
0277   for (std::size_t seed_n = 0; seed_n < max_seed_n; seed_n++) {
0278     const auto& seed_features = features.seed_features[seed_n];
0279     ptr = &data_[kSeedingTracks][offset + seed_n * n_features_seed_];
0280     start = ptr;
0281     seedTrack_tensor_filler(ptr, seed_features);
0282     assert(start + n_features_seed_ - 1 == ptr);
0283   }
0284
0285   // neighbours
0286   offset = i_jet * input_sizes_[kNeighbourTracks];
0287   for (std::size_t seed_n = 0; seed_n < max_seed_n; seed_n++) {
0288     const auto& neighbourTracks_features = features.seed_features[seed_n].nearTracks;
0289     auto max_neighbour_n = std::min(neighbourTracks_features.size(), (std::size_t)n_neighbor_);
0290     for (std::size_t neighbour_n = 0; neighbour_n < max_neighbour_n; neighbour_n++) {
0291       ptr = &data_[kNeighbourTracks + seed_n][offset + neighbour_n * n_features_neighbor_];
0292       start = ptr;
0293       neighbourTrack_tensor_filler(ptr, neighbourTracks_features[neighbour_n]);
0294       assert(start + n_features_neighbor_ - 1 == ptr);
0295     }
0296   }
0297 }
0298
0299 //define this as a plug-in
0300 DEFINE_FWK_MODULE(DeepCombinedONNXJetTagsProducer);