ONNXRuntime/plugins/DeepVertexONNXJetTagsProducer.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 DeepVertexONNXJetTagsProducer : public edm::stream::EDProducer<edm::GlobalCache<ONNXRuntime>> {
0024 public:
0025   explicit DeepVertexONNXJetTagsProducer(const edm::ParameterSet&, const ONNXRuntime*);
0026   ~DeepVertexONNXJetTagsProducer() 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 { kGlobal = 0, kSeedingTracks = 1, kNeighbourTracks = 2 };
0050   const static unsigned n_features_global_ = deepvertex::n_features_global;
0051   const static unsigned n_seed_ = deepvertex::n_seed;
0052   const static unsigned n_features_seed_ = deepvertex::n_features_seed;
0053   const static unsigned n_neighbor_ = deepvertex::n_neighbor;
0054   const static unsigned n_features_neighbor_ = deepvertex::n_features_neighbor;
0055
0056   const static std::vector<unsigned> input_sizes_;
0057
0058   // hold the input data
0059   FloatArrays data_;
0060 };
0061
0062 const std::vector<unsigned> DeepVertexONNXJetTagsProducer::input_sizes_{n_features_global_,
0063                                                                         n_seed_* n_features_seed_,
0064                                                                         n_neighbor_* n_features_neighbor_,
0065                                                                         n_neighbor_* n_features_neighbor_,
0066                                                                         n_neighbor_* n_features_neighbor_,
0067                                                                         n_neighbor_* n_features_neighbor_,
0068                                                                         n_neighbor_* n_features_neighbor_,
0069                                                                         n_neighbor_* n_features_neighbor_,
0070                                                                         n_neighbor_* n_features_neighbor_,
0071                                                                         n_neighbor_* n_features_neighbor_,
0072                                                                         n_neighbor_* n_features_neighbor_,
0073                                                                         n_neighbor_* n_features_neighbor_};
0074
0075 DeepVertexONNXJetTagsProducer::DeepVertexONNXJetTagsProducer(const edm::ParameterSet& iConfig, const ONNXRuntime* cache)
0076     : src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
0077       flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
0078       input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
0079       output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")),
0080       min_jet_pt_(iConfig.getParameter<double>("min_jet_pt")),
0081       max_jet_eta_(iConfig.getParameter<double>("max_jet_eta")) {
0082   // get output names from flav_names
0083   for (const auto& flav_name : flav_names_) {
0084     produces<JetTagCollection>(flav_name);
0085   }
0086
0087   assert(input_names_.size() == input_sizes_.size());
0088 }
0089
0090 DeepVertexONNXJetTagsProducer::~DeepVertexONNXJetTagsProducer() {}
0091
0092 void DeepVertexONNXJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0093   // pfDeepFlavourJetTags
0094   edm::ParameterSetDescription desc;
0095   desc.add<edm::InputTag>("src", edm::InputTag("pfDeepFlavourTagInfos"));
0096   desc.add<std::vector<std::string>>("input_names",
0097                                      {"input_1",
0098                                       "input_2",
0099                                       "input_3",
0100                                       "input_4",
0101                                       "input_5",
0102                                       "input_6",
0103                                       "input_7",
0104                                       "input_8",
0105                                       "input_9",
0106                                       "input_10",
0107                                       "input_11",
0108                                       "input_12"});
0109   desc.add<edm::FileInPath>("model_path", edm::FileInPath("RecoBTag/Combined/data/DeepVertex/phase1_deepvertex.onnx"));
0110   desc.add<std::vector<std::string>>("output_names", {"dense_6"});
0111   desc.add<std::vector<std::string>>("flav_names", std::vector<std::string>{"probb", "probc", "probuds", "probg"});
0112   desc.add<double>("min_jet_pt", 15.0);
0113   desc.add<double>("max_jet_eta", 2.5);
0114
0115   descriptions.add("pfDeepVertexJetTags", desc);
0116 }
0117
0118 std::unique_ptr<ONNXRuntime> DeepVertexONNXJetTagsProducer::initializeGlobalCache(const edm::ParameterSet& iConfig) {
0119   return std::make_unique<ONNXRuntime>(iConfig.getParameter<edm::FileInPath>("model_path").fullPath());
0120 }
0121
0122 void DeepVertexONNXJetTagsProducer::globalEndJob(const ONNXRuntime* cache) {}
0123
0124 void DeepVertexONNXJetTagsProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0125   edm::Handle<TagInfoCollection> tag_infos;
0126   iEvent.getByToken(src_, tag_infos);
0127
0128   data_.clear();
0129
0130   std::vector<std::unique_ptr<JetTagCollection>> output_tags;
0131   if (!tag_infos->empty()) {
0132     unsigned good_taginfo_count = 0;
0133     std::vector<bool> good_taginfo_jets(tag_infos->size(), false);
0134     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0135       const auto& jet_ref = (*tag_infos)[jet_n].jet();
0136       if (jet_ref->pt() > min_jet_pt_ && std::fabs(jet_ref->eta()) < max_jet_eta_) {
0137         good_taginfo_count++;
0138         good_taginfo_jets[jet_n] = true;
0139       }
0140     }
0141
0142     // init data storage w correct size
0143     for (const auto& len : input_sizes_) {
0144       data_.emplace_back(good_taginfo_count * len, 0);
0145     }
0146
0147     // initialize output collection
0148     auto jet_ref = tag_infos->begin()->jet();
0149     auto ref2prod = edm::makeRefToBaseProdFrom(jet_ref, iEvent);
0150     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0151       output_tags.emplace_back(std::make_unique<JetTagCollection>(ref2prod));
0152     }
0153
0154     // convert inputs
0155     unsigned inputs_done_count = 0;
0156     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0157       if (good_taginfo_jets[jet_n]) {
0158         const auto& taginfo = (*tag_infos)[jet_n];
0159         make_inputs(inputs_done_count, taginfo);
0160         inputs_done_count++;
0161       }
0162     }
0163
0164     // run prediction
0165     assert(inputs_done_count == good_taginfo_count);
0166     const auto outputs = globalCache()->run(input_names_, data_, {}, output_names_, good_taginfo_count)[0];
0167     assert(outputs.size() == flav_names_.size() * good_taginfo_count);
0168
0169     // get the outputs
0170     unsigned i_output = 0;
0171     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0172       const auto& jet_ref = (*tag_infos)[jet_n].jet();
0173       for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0174         if (good_taginfo_jets[jet_n]) {
0175           (*(output_tags[flav_n]))[jet_ref] = outputs[i_output];
0176           ++i_output;
0177         } else {
0178           (*(output_tags[flav_n]))[jet_ref] = -2;
0179         }
0180       }
0181     }
0182   } else {
0183     // create empty output collection
0184     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0185       output_tags.emplace_back(std::make_unique<JetTagCollection>());
0186     }
0187   }
0188
0189   // put into the event
0190   for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) {
0191     iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]);
0192   }
0193 }
0194
0195 void DeepVertexONNXJetTagsProducer::make_inputs(unsigned i_jet, const reco::DeepFlavourTagInfo& taginfo) {
0196   const auto& features = taginfo.features();
0197   float* ptr = nullptr;
0198   const float* start = nullptr;
0199   unsigned offset = 0;
0200
0201   // jet variables
0202   offset = i_jet * input_sizes_[kGlobal];
0203   const auto& jet_features = features.jet_features;
0204   ptr = &data_[kGlobal][offset];
0205   start = ptr;
0206   jet4vec_tensor_filler(ptr, jet_features);
0207   assert(start + n_features_global_ - 1 == ptr);
0208
0209   // seeds
0210   auto max_seed_n = std::min(features.seed_features.size(), (std::size_t)n_seed_);
0211   offset = i_jet * input_sizes_[kSeedingTracks];
0212   for (std::size_t seed_n = 0; seed_n < max_seed_n; seed_n++) {
0213     const auto& seed_features = features.seed_features[seed_n];
0214     ptr = &data_[kSeedingTracks][offset + seed_n * n_features_seed_];
0215     start = ptr;
0216     seedTrack_tensor_filler(ptr, seed_features);
0217     assert(start + n_features_seed_ - 1 == ptr);
0218   }
0219
0220   // neighbours
0221   offset = i_jet * input_sizes_[kNeighbourTracks];
0222   for (std::size_t seed_n = 0; seed_n < max_seed_n; seed_n++) {
0223     const auto& neighbourTracks_features = features.seed_features[seed_n].nearTracks;
0224     auto max_neighbour_n = std::min(neighbourTracks_features.size(), (std::size_t)n_neighbor_);
0225     for (std::size_t neighbour_n = 0; neighbour_n < max_neighbour_n; neighbour_n++) {
0226       ptr = &data_[kNeighbourTracks + seed_n][offset + neighbour_n * n_features_neighbor_];
0227       start = ptr;
0228       neighbourTrack_tensor_filler(ptr, neighbourTracks_features[neighbour_n]);
0229       assert(start + n_features_neighbor_ - 1 == ptr);
0230     }
0231   }
0232 }
0233
0234 //define this as a plug-in
0235 DEFINE_FWK_MODULE(DeepVertexONNXJetTagsProducer);