ONNXRuntime/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.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/UnifiedParticleTransformerAK4TagInfo.h"
0015 #include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4Features.h"
0016
0017 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0018
0019 #include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h"
0020 #include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
0021
0022 using namespace cms::Ort;
0023
0024 class UnifiedParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer<edm::GlobalCache<ONNXRuntime>> {
0025 public:
0026   explicit UnifiedParticleTransformerAK4ONNXJetTagsProducer(const edm::ParameterSet&, const ONNXRuntime*);
0027   ~UnifiedParticleTransformerAK4ONNXJetTagsProducer() override = default;
0028
0029   static void fillDescriptions(edm::ConfigurationDescriptions&);
0030
0031   static std::unique_ptr<ONNXRuntime> initializeGlobalCache(const edm::ParameterSet&);
0032   static void globalEndJob(const ONNXRuntime*);
0033
0034 private:
0035   typedef std::vector<reco::UnifiedParticleTransformerAK4TagInfo> TagInfoCollection;
0036   typedef reco::JetTagCollection JetTagCollection;
0037
0038   void produce(edm::Event&, const edm::EventSetup&) override;
0039
0040   void make_inputs(btagbtvdeep::UnifiedParticleTransformerAK4Features features);
0041   void get_input_sizes(const reco::FeaturesTagInfo<btagbtvdeep::UnifiedParticleTransformerAK4Features> taginfo);
0042
0043   const edm::EDGetTokenT<TagInfoCollection> src_;
0044   std::vector<std::string> flav_names_;
0045   std::vector<std::string> input_names_;
0046   bool use_dynamic_axes_ = false;
0047   std::vector<std::string> output_names_;
0048
0049   unsigned n_cpf_;
0050   unsigned n_lt_;
0051   unsigned n_npf_;
0052   unsigned n_sv_;
0053   std::vector<unsigned> input_sizes_;
0054   std::vector<std::vector<int64_t>> input_shapes_;  // shapes of each input group (-1 for dynamic axis)
0055
0056   // hold the input data
0057   FloatArrays data_;
0058 };
0059
0060 UnifiedParticleTransformerAK4ONNXJetTagsProducer::UnifiedParticleTransformerAK4ONNXJetTagsProducer(
0061     const edm::ParameterSet& iConfig, const ONNXRuntime* cache)
0062     : src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
0063       flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
0064       input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
0065       use_dynamic_axes_(iConfig.getParameter<edm::FileInPath>("model_path").fullPath().find("V01") !=
0066                         std::string::npos),
0067       output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")) {
0068   // get output names from flav_names
0069   for (const auto& flav_name : flav_names_) {
0070     produces<JetTagCollection>(flav_name);
0071   }
0072 }
0073
0074 void UnifiedParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0075   // pfUnifiedParticleTransformerAK4JetTags
0076   edm::ParameterSetDescription desc;
0077   desc.add<edm::InputTag>("src", edm::InputTag("pfUnifiedParticleTransformerAK4TagInfos"));
0078   desc.add<std::vector<std::string>>(
0079       "input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6", "input_7", "input_8"});
0080   desc.add<edm::FileInPath>("model_path",
0081                             edm::FileInPath("RecoBTag/Combined/data/UParTAK4/PUPPI/V01/modelfile/model.onnx"));
0082   desc.add<std::vector<std::string>>("output_names", {"softmax"});
0083   desc.add<std::vector<std::string>>(
0084       "flav_names",
0085       std::vector<std::string>{"probb",        "probbb",       "problepb",     "probc",         "probs",
0086                                "probu",        "probd",        "probg",        "probele",       "probmu",
0087                                "probtaup1h0p", "probtaup1h1p", "probtaup1h2p", "probtaup3h0p",  "probtaup3h1p",
0088                                "probtaum1h0p", "probtaum1h1p", "probtaum1h2p", "probtaum3h0p",  "probtaum3h1p",
0089                                "ptcorr",       "ptreshigh",    "ptreslow",     "ptnu",          "probemudata",
0090                                "probemumc",    "probdimudata", "probdimumc",   "probmutaudata", "probmutaumc"});
0091
0092   descriptions.add("pfUnifiedParticleTransformerAK4JetTags", desc);
0093 }
0094
0095 std::unique_ptr<ONNXRuntime> UnifiedParticleTransformerAK4ONNXJetTagsProducer::initializeGlobalCache(
0096     const edm::ParameterSet& iConfig) {
0097   return std::make_unique<ONNXRuntime>(iConfig.getParameter<edm::FileInPath>("model_path").fullPath());
0098 }
0099
0100 void UnifiedParticleTransformerAK4ONNXJetTagsProducer::globalEndJob(const ONNXRuntime* cache) {}
0101
0102 void UnifiedParticleTransformerAK4ONNXJetTagsProducer::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)UparT::N_InputFeatures.at(UparT::kChargedCandidates)},
0128                        {(int64_t)1, (int64_t)n_lt_, (int64_t)UparT::N_InputFeatures.at(UparT::kLostTracks)},
0129                        {(int64_t)1, (int64_t)n_npf_, (int64_t)UparT::N_InputFeatures.at(UparT::kNeutralCandidates)},
0130                        {(int64_t)1, (int64_t)n_sv_, (int64_t)UparT::N_InputFeatures.at(UparT::kVertices)},
0131                        {(int64_t)1, (int64_t)n_cpf_, (int64_t)UparT::N_InputFeatures.at(UparT::kChargedCandidates4Vec)},
0132                        {(int64_t)1, (int64_t)n_lt_, (int64_t)UparT::N_InputFeatures.at(UparT::kLostTracks4Vec)},
0133                        {(int64_t)1, (int64_t)n_npf_, (int64_t)UparT::N_InputFeatures.at(UparT::kNeutralCandidates4Vec)},
0134                        {(int64_t)1, (int64_t)n_sv_, (int64_t)UparT::N_InputFeatures.at(UparT::kVertices4Vec)}};
0135
0136       outputs = globalCache()->run(input_names_, data_, input_shapes_, output_names_, 1)[0];
0137       assert(outputs.size() == flav_names_.size());
0138     }
0139
0140     const auto& jet_ref = tag_infos->at(jet_n).jet();
0141     for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0142       (*(output_tags[flav_n]))[jet_ref] = outputs[flav_n];
0143     }
0144   }
0145
0146   // put into the event
0147   for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) {
0148     iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]);
0149   }
0150 }
0151
0152 void UnifiedParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
0153     const reco::FeaturesTagInfo<btagbtvdeep::UnifiedParticleTransformerAK4Features> taginfo) {
0154   const auto& features = taginfo.features();
0155
0156   if (use_dynamic_axes_) {
0157     // Use actual sizes for dynamic axes version
0158     n_cpf_ =
0159         std::clamp((unsigned int)features.c_pf_features.size(), (unsigned int)1, (unsigned int)UparT::n_cpf_accept);
0160     n_lt_ = std::clamp((unsigned int)features.lt_features.size(), (unsigned int)1, (unsigned int)UparT::n_lt_accept);
0161     n_npf_ =
0162         std::clamp((unsigned int)features.n_pf_features.size(), (unsigned int)1, (unsigned int)UparT::n_npf_accept);
0163     n_sv_ = std::clamp((unsigned int)features.sv_features.size(), (unsigned int)1, (unsigned int)UparT::n_sv_accept);
0164
0165   } else {
0166     // Use fixed sizes for original version
0167     n_cpf_ = (unsigned int)UparT::n_cpf_accept;
0168     n_lt_ = (unsigned int)UparT::n_lt_accept;
0169     n_npf_ = (unsigned int)UparT::n_npf_accept;
0170     n_sv_ = (unsigned int)UparT::n_sv_accept;
0171   }
0172
0173   input_sizes_ = {
0174       n_cpf_ * UparT::N_InputFeatures.at(UparT::kChargedCandidates),
0175       n_lt_ * UparT::N_InputFeatures.at(UparT::kLostTracks),
0176       n_npf_ * UparT::N_InputFeatures.at(UparT::kNeutralCandidates),
0177       n_sv_ * UparT::N_InputFeatures.at(UparT::kVertices),
0178       n_cpf_ * UparT::N_InputFeatures.at(UparT::kChargedCandidates4Vec),
0179       n_lt_ * UparT::N_InputFeatures.at(UparT::kLostTracks4Vec),
0180       n_npf_ * UparT::N_InputFeatures.at(UparT::kNeutralCandidates4Vec),
0181       n_sv_ * UparT::N_InputFeatures.at(UparT::kVertices4Vec),
0182   };
0183   // init data storage
0184   data_.clear();
0185   for (const auto& len : input_sizes_) {
0186     data_.emplace_back(1 * len, 0);
0187   }
0188
0189   make_inputs(features);
0190 }
0191
0192 void UnifiedParticleTransformerAK4ONNXJetTagsProducer::make_inputs(
0193     btagbtvdeep::UnifiedParticleTransformerAK4Features features) {
0194   const float* start = nullptr;
0195   unsigned offset = 0;
0196
0197   auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
0198   auto max_lt_n = std::min(features.lt_features.size(), (std::size_t)n_lt_);
0199   auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
0200   auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_);
0201   // c_pf candidates
0202   UparT_tensor_filler(data_, UparT::kChargedCandidates, features.c_pf_features, max_c_pf_n, start, offset);
0203   // lt candidates
0204   UparT_tensor_filler(data_, UparT::kLostTracks, features.lt_features, max_lt_n, start, offset);
0205   // n_pf candidates
0206   UparT_tensor_filler(data_, UparT::kNeutralCandidates, features.n_pf_features, max_n_pf_n, start, offset);
0207   // sv candidates
0208   UparT_tensor_filler(data_, UparT::kVertices, features.sv_features, max_sv_n, start, offset);
0209   // cpf pairwise features (4-vectors)
0210   UparT_tensor_filler(data_, UparT::kChargedCandidates4Vec, features.c_pf_features, max_c_pf_n, start, offset);
0211   // lt pairwise features (4-vectors) specific case requiring (pt,eta,phi,e)
0212   UparT_tensor_filler(data_, UparT::kLostTracks4Vec, features.lt_features, max_lt_n, start, offset);
0213   // npf pairwise features (4-vectors)
0214   UparT_tensor_filler(data_, UparT::kNeutralCandidates4Vec, features.n_pf_features, max_n_pf_n, start, offset);
0215   // sv pairwise features (4-vectors)
0216   UparT_tensor_filler(data_, UparT::kVertices4Vec, features.sv_features, max_sv_n, start, offset);
0217 }
0218
0219 //define this as a plug-in
0220 DEFINE_FWK_MODULE(UnifiedParticleTransformerAK4ONNXJetTagsProducer);