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 #include "RecoBTag/ONNXRuntime/interface/tensor_fillers.h"
0020 #include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
0021
0022 using namespace cms::Ort;
0023
0024 class ParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDProducer<edm::GlobalCache<ONNXRuntime>> {
0025 public:
0026   explicit ParticleTransformerAK4ONNXJetTagsProducer(const edm::ParameterSet&, const ONNXRuntime*);
0027   ~ParticleTransformerAK4ONNXJetTagsProducer() 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::ParticleTransformerAK4TagInfo> TagInfoCollection;
0036   typedef reco::JetTagCollection JetTagCollection;
0037
0038   void produce(edm::Event&, const edm::EventSetup&) override;
0039
0040   void make_inputs(btagbtvdeep::ParticleTransformerAK4Features features);
0041   void get_input_sizes(const reco::FeaturesTagInfo<btagbtvdeep::ParticleTransformerAK4Features> taginfo);
0042
0043   const edm::EDGetTokenT<TagInfoCollection> src_;
0044   std::vector<std::string> flav_names_;
0045   std::vector<std::string> input_names_;
0046   std::vector<std::string> output_names_;
0047   unsigned int n_cpf_;
0048   unsigned int n_npf_;
0049   unsigned int n_sv_;
0050   std::vector<unsigned> input_sizes_;
0051   std::vector<std::vector<int64_t>> input_shapes_;  // shapes of each input group (-1 for dynamic axis)
0052
0053   // hold the input data
0054   FloatArrays data_;
0055 };
0056
0057 ParticleTransformerAK4ONNXJetTagsProducer::ParticleTransformerAK4ONNXJetTagsProducer(const edm::ParameterSet& iConfig,
0058                                                                                      const ONNXRuntime* cache)
0059     : src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
0060       flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
0061       input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
0062       output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")) {
0063   // get output names from flav_names
0064   for (const auto& flav_name : flav_names_) {
0065     produces<JetTagCollection>(flav_name);
0066   }
0067 }
0068
0069 void ParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0070   // pfParticleTransformerAK4JetTags
0071   edm::ParameterSetDescription desc;
0072   desc.add<edm::InputTag>("src", edm::InputTag("pfParticleTransformerAK4TagInfos"));
0073   desc.add<std::vector<std::string>>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6"});
0074   desc.add<edm::FileInPath>("model_path",
0075                             edm::FileInPath("RecoBTag/Combined/data/RobustParTAK4/PUPPI/V00/modelfile/model.onnx"));
0076   desc.add<std::vector<std::string>>("output_names", {"softmax"});
0077   desc.add<std::vector<std::string>>(
0078       "flav_names", std::vector<std::string>{"probb", "probbb", "problepb", "probc", "probuds", "probg"});
0079
0080   descriptions.add("pfParticleTransformerAK4JetTags", desc);
0081 }
0082
0083 std::unique_ptr<ONNXRuntime> ParticleTransformerAK4ONNXJetTagsProducer::initializeGlobalCache(
0084     const edm::ParameterSet& iConfig) {
0085   return std::make_unique<ONNXRuntime>(iConfig.getParameter<edm::FileInPath>("model_path").fullPath());
0086 }
0087
0088 void ParticleTransformerAK4ONNXJetTagsProducer::globalEndJob(const ONNXRuntime* cache) {}
0089
0090 void ParticleTransformerAK4ONNXJetTagsProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0091   edm::Handle<TagInfoCollection> tag_infos;
0092   iEvent.getByToken(src_, tag_infos);
0093
0094   // initialize output collection
0095   std::vector<std::unique_ptr<JetTagCollection>> output_tags;
0096   if (!tag_infos->empty()) {
0097     auto jet_ref = tag_infos->begin()->jet();
0098     auto ref2prod = edm::makeRefToBaseProdFrom(jet_ref, iEvent);
0099     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0100       output_tags.emplace_back(std::make_unique<JetTagCollection>(ref2prod));
0101     }
0102   } else {
0103     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0104       output_tags.emplace_back(std::make_unique<JetTagCollection>());
0105     }
0106   }
0107
0108   for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0109     const auto& taginfo = (*tag_infos)[jet_n];
0110     std::vector<float> outputs(flav_names_.size(), -1.0);
0111     if (taginfo.features().is_filled) {
0112       get_input_sizes(taginfo);
0113
0114       // run prediction with dynamic batch size per event
0115       input_shapes_ = {{(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::N_InputFeatures.at(parT::kChargedCandidates)},
0116                        {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::N_InputFeatures.at(parT::kNeutralCandidates)},
0117                        {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::N_InputFeatures.at(parT::kVertices)},
0118                        {(int64_t)1, (int64_t)n_cpf_, (int64_t)parT::N_InputFeatures.at(parT::kChargedCandidates4Vec)},
0119                        {(int64_t)1, (int64_t)n_npf_, (int64_t)parT::N_InputFeatures.at(parT::kNeutralCandidates4Vec)},
0120                        {(int64_t)1, (int64_t)n_sv_, (int64_t)parT::N_InputFeatures.at(parT::kVertices4Vec)}};
0121
0122       outputs = globalCache()->run(input_names_, data_, input_shapes_, output_names_, 1)[0];
0123       assert(outputs.size() == flav_names_.size());
0124     }
0125
0126     const auto& jet_ref = tag_infos->at(jet_n).jet();
0127     for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0128       (*(output_tags[flav_n]))[jet_ref] = outputs[flav_n];
0129     }
0130   }
0131
0132   // put into the event
0133   for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) {
0134     iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]);
0135   }
0136 }
0137
0138 void ParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
0139     const reco::FeaturesTagInfo<btagbtvdeep::ParticleTransformerAK4Features> taginfo) {
0140   const auto& features = taginfo.features();
0141
0142   n_cpf_ = std::clamp((unsigned int)features.c_pf_features.size(), (unsigned int)1, (unsigned int)parT::n_cpf_accept);
0143   n_npf_ = std::clamp((unsigned int)features.n_pf_features.size(), (unsigned int)1, (unsigned int)parT::n_npf_accept);
0144   n_sv_ = std::clamp((unsigned int)features.sv_features.size(), (unsigned int)1, (unsigned int)parT::n_sv_accept);
0145
0146   input_sizes_ = {
0147       n_cpf_ * parT::N_InputFeatures.at(parT::kChargedCandidates),
0148       n_npf_ * parT::N_InputFeatures.at(parT::kNeutralCandidates),
0149       n_sv_ * parT::N_InputFeatures.at(parT::kVertices),
0150       n_cpf_ * parT::N_InputFeatures.at(parT::kChargedCandidates4Vec),
0151       n_npf_ * parT::N_InputFeatures.at(parT::kNeutralCandidates4Vec),
0152       n_sv_ * parT::N_InputFeatures.at(parT::kVertices4Vec),
0153   };
0154   // init data storage
0155   data_.clear();
0156   for (const auto& len : input_sizes_) {
0157     data_.emplace_back(1 * len, 0);
0158   }
0159
0160   make_inputs(features);
0161 }
0162
0163 void ParticleTransformerAK4ONNXJetTagsProducer::make_inputs(btagbtvdeep::ParticleTransformerAK4Features features) {
0164   //float* ptr = nullptr;
0165   const float* start = nullptr;
0166   unsigned offset = 0;
0167
0168   auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)n_cpf_);
0169   auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)n_npf_);
0170   auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)n_sv_);
0171
0172   // c_pf candidates
0173   parT_tensor_filler(data_, parT::kChargedCandidates, features.c_pf_features, max_c_pf_n, start, offset);
0174   // n_pf candidates
0175   parT_tensor_filler(data_, parT::kNeutralCandidates, features.n_pf_features, max_n_pf_n, start, offset);
0176   // sv candidates
0177   parT_tensor_filler(data_, parT::kVertices, features.sv_features, max_sv_n, start, offset);
0178   // cpf pairwise features (4-vectors)
0179   parT_tensor_filler(data_, parT::kChargedCandidates4Vec, features.c_pf_features, max_c_pf_n, start, offset);
0180   // npf pairwise features (4-vectors)
0181   parT_tensor_filler(data_, parT::kNeutralCandidates4Vec, features.n_pf_features, max_n_pf_n, start, offset);
0182   // sv pairwise features (4-vectors)
0183   parT_tensor_filler(data_, parT::kVertices4Vec, features.sv_features, max_sv_n, start, offset);
0184 }
0185
0186 //define this as a plug-in
0187 DEFINE_FWK_MODULE(ParticleTransformerAK4ONNXJetTagsProducer);