HcalTasks/plugins/OnlineDQMDigiAD_cmssw.cc

0001 /*
0002  * OnlineDQMDigiAD_cmssw.cpp
0003  *
0004  * Created on: Jun 10, 2023
0005  * Author: Mulugeta W.Asres, UiA, Norway
0006  *
0007  * The implementation follows https://github.com/cms-sw/cmssw/tree/master/PhysicsTools/ONNXRuntime
0008  */
0009
0010 // #include "FWCore/Utilities/interface/Exception.h"
0011 // #include "FWCore/Utilities/interface/thread_safety_macros.h"
0012 // #include "FWCore/Framework/interface/Event.h"
0013 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0014 #include "FWCore/ParameterSet/interface/FileInPath.h"
0015 #include "HeterogeneousCore/CUDAUtilities/interface/requireDevices.h"
0016
0017 #include <algorithm>
0018 #include <cassert>
0019 #include <functional>
0020 #include <iostream>
0021 #include <memory>
0022 #include <numeric>
0023 #include <algorithm>
0024
0025 #include "DQM/HcalTasks/plugins/OnlineDQMDigiAD_cmssw.h"
0026
0027 // using namespace std;
0028 using namespace cms::Ort;
0029
0030 // Constructor
0031 OnlineDQMDigiAD::OnlineDQMDigiAD(const std::string model_system_name,
0032                                  const std::string &modelFilepath,
0033                                  Backend backend) {
0034   std::string instanceName{"DESMOD Digioccupancy Map AD inference"};
0035
0036   /**************** Initailize Model Memory States ******************/
0037   InitializeState();  // initailize model memory states to zero
0038
0039   /**************** Create ORT session ******************/
0040   // Set up options for session
0041   auto session_options = ONNXRuntime::defaultSessionOptions(backend);
0042   // Create session by loading the onnx model
0043   model_path = edm::FileInPath(modelFilepath).fullPath();
0044   auto uOrtSession = std::make_unique<ONNXRuntime>(model_path, &session_options);
0045   ort_mSession = std::move(uOrtSession);
0046
0047   // check model availability
0048   hcal_subsystem_name = model_system_name;
0049
0050   IsModelExist(hcal_subsystem_name);  // assert model integration for the given hcal system name
0051
0052   if (hcal_subsystem_name == "he") {
0053     std::vector<std::vector<int64_t>> input_shapes_ = {
0054         {batch_size, 64, 72, 7, 1},
0055         {batch_size, 1},
0056         {1, 1},
0057         {batch_size, model_state_inner_dim, model_state_layer_dims[0][0]},
0058         {batch_size, model_state_inner_dim, model_state_layer_dims[0][0]},
0059         {batch_size, model_state_inner_dim, model_state_layer_dims[0][1]},
0060         {batch_size, model_state_inner_dim, model_state_layer_dims[0][1]},
0061         {batch_size, model_state_inner_dim, model_state_layer_dims[1][0]},
0062         {batch_size, model_state_inner_dim, model_state_layer_dims[1][0]},
0063         {batch_size, model_state_inner_dim, model_state_layer_dims[1][1]},
0064         {batch_size, model_state_inner_dim, model_state_layer_dims[1][1]}};  // input dims
0065     input_shapes = input_shapes_;
0066   }
0067
0068   else if (hcal_subsystem_name == "hb") {
0069     std::vector<std::vector<int64_t>> input_shapes_ = {
0070         {batch_size, 64, 72, 4, 1},
0071         {batch_size, 1},
0072         {1, 1},
0073         {batch_size, model_state_inner_dim, model_state_layer_dims[0][0]},
0074         {batch_size, model_state_inner_dim, model_state_layer_dims[0][0]},
0075         {batch_size, model_state_inner_dim, model_state_layer_dims[0][1]},
0076         {batch_size, model_state_inner_dim, model_state_layer_dims[0][1]},
0077         {batch_size, model_state_inner_dim, model_state_layer_dims[1][0]},
0078         {batch_size, model_state_inner_dim, model_state_layer_dims[1][0]},
0079         {batch_size, model_state_inner_dim, model_state_layer_dims[1][1]},
0080         {batch_size, model_state_inner_dim, model_state_layer_dims[1][1]}};  // input dims
0081     input_shapes = input_shapes_;
0082   }
0083 }
0084
0085 void OnlineDQMDigiAD::IsModelExist(std::string hcal_subsystem_name) {
0086   if (std::find(hcal_modeled_systems.begin(), hcal_modeled_systems.end(), hcal_subsystem_name) ==
0087       hcal_modeled_systems.end()) {
0088     std::string err =
0089         "ML for OnlineDQM is not currently supported for the selected " + hcal_subsystem_name + " system!\n";
0090     throw std::invalid_argument(err);
0091   }
0092 }
0093
0094 void OnlineDQMDigiAD::InitializeState() {
0095   // model memory states vectors init, only when the runs starts or for the first LS
0096   std::fill(input_model_state_memory_e_0_0.begin(),
0097             input_model_state_memory_e_0_0.end(),
0098             float(0.0));  // init model memory states-encoder_layer_0_state_0 to zero
0099   std::fill(input_model_state_memory_e_0_1.begin(),
0100             input_model_state_memory_e_0_1.end(),
0101             float(0.0));  // init model memory states-encoder_layer_0_state_1 to zero
0102   std::fill(input_model_state_memory_e_1_0.begin(),
0103             input_model_state_memory_e_1_0.end(),
0104             float(0.0));  // init model memory states-encoder_layer_1_state_0 to zero
0105   std::fill(input_model_state_memory_e_1_1.begin(),
0106             input_model_state_memory_e_1_1.end(),
0107             float(0.0));  // init model memory states-encoder_layer_1_state_1 to zero
0108   std::fill(input_model_state_memory_d_0_0.begin(),
0109             input_model_state_memory_d_0_0.end(),
0110             float(0.0));  // init model memory states-decoder_layer_0_state_0 to zero
0111   std::fill(input_model_state_memory_d_0_1.begin(),
0112             input_model_state_memory_d_0_1.end(),
0113             float(0.0));  // init model memory states-decoder_layer_0_state_1 to zero
0114   std::fill(input_model_state_memory_d_1_0.begin(),
0115             input_model_state_memory_d_1_0.end(),
0116             float(0.0));  // init model memory states-decoder_layer_1_state_0 to zero
0117   std::fill(input_model_state_memory_d_1_1.begin(),
0118             input_model_state_memory_d_1_1.end(),
0119             float(0.0));  // init model memory states-decoder_layer_1_state_1 to zero
0120
0121   // model_state_refresh_counter = 15; // counter set due to onnx double datatype handling limitation that might cause precision error to propagate.
0122   model_state_refresh_counter =
0123       1;  // DQM multithread returns non-sequential LS. Hence, the model will not keep states (experimental)
0124 }
0125
0126 std::vector<float> OnlineDQMDigiAD::Serialize2DVector(const std::vector<std::vector<float>> &input_2d_vec) {
0127   std::vector<float> output;
0128   for (const auto &row : input_2d_vec) {
0129     for (const auto &element : row) {
0130       output.push_back(element);
0131     }
0132   }
0133   return output;
0134 }
0135
0136 std::vector<std::vector<float>> OnlineDQMDigiAD::Map1DTo2DVector(const std::vector<float> &input_1d_vec,
0137                                                                  const int numSplits) {
0138   if (numSplits <= 0)
0139     throw std::invalid_argument("numSplits must be greater than 0.");
0140
0141   std::size_t const splitted_size = input_1d_vec.size() / numSplits;
0142
0143   if (splitted_size * numSplits != input_1d_vec.size())
0144     throw std::invalid_argument("Conversion is not allowed! The input vector length " +
0145                                 std::to_string(input_1d_vec.size()) + " must be divisible by the numSplits " +
0146                                 std::to_string(numSplits) + ".");
0147
0148   std::vector<std::vector<float>> output_2d_vec;
0149
0150   for (int i = 0; i < numSplits; i++) {
0151     std::vector<float> chunch_vec(input_1d_vec.begin() + i * splitted_size,
0152                                   input_1d_vec.begin() + (i + 1) * splitted_size);
0153     output_2d_vec.push_back(chunch_vec);
0154   }
0155   return output_2d_vec;
0156 }
0157
0158 std::vector<float> OnlineDQMDigiAD::PrepareONNXDQMMapVectors(
0159     std::vector<std::vector<std::vector<float>>> &digiHcal2DHist_depth_all) {
0160   std::vector<float> digi3DHistVector_serialized;
0161
0162   for (const std::vector<std::vector<float>> &digiHcal2DHist_depth : digiHcal2DHist_depth_all) {
0163     std::vector<float> digiHcalDHist_serialized_depth = Serialize2DVector(digiHcal2DHist_depth);
0164     digi3DHistVector_serialized.insert(digi3DHistVector_serialized.end(),
0165                                        digiHcalDHist_serialized_depth.begin(),
0166                                        digiHcalDHist_serialized_depth.end());
0167   }
0168
0169   return digi3DHistVector_serialized;
0170 }
0171
0172 std::vector<std::vector<std::vector<float>>> OnlineDQMDigiAD::ONNXOutputToDQMHistMap(
0173     const std::vector<std::vector<float>> &ad_model_output_vectors,
0174     const int numDepth,
0175     const int numDIeta,
0176     const int selOutputIdx) {
0177   // each output_vector is a serialized 3d hist map
0178
0179   const std::vector<float> &output_vector = ad_model_output_vectors[selOutputIdx];
0180   std::vector<std::vector<float>> output_2d_vec = Map1DTo2DVector(output_vector, numDepth);
0181
0182   std::vector<std::vector<std::vector<float>>> digiHcal3DHist;
0183   for (const std::vector<float> &output_vector_depth : output_2d_vec) {
0184     std::vector<std::vector<float>> digiHcal2DHist_depth = Map1DTo2DVector(output_vector_depth, numDIeta);
0185     digiHcal3DHist.push_back(digiHcal2DHist_depth);
0186   }
0187
0188   return digiHcal3DHist;
0189 }
0190
0191 // Perform inference for a given dqm map
0192 std::vector<std::vector<float>> OnlineDQMDigiAD::Inference(std::vector<float> &digiHcalMapTW,
0193                                                            std::vector<float> &numEvents,
0194                                                            std::vector<float> &adThr,
0195                                                            std::vector<float> &input_model_state_memory_e_0_0,
0196                                                            std::vector<float> &input_model_state_memory_e_0_1,
0197                                                            std::vector<float> &input_model_state_memory_e_1_0,
0198                                                            std::vector<float> &input_model_state_memory_e_1_1,
0199                                                            std::vector<float> &input_model_state_memory_d_0_0,
0200                                                            std::vector<float> &input_model_state_memory_d_0_1,
0201                                                            std::vector<float> &input_model_state_memory_d_1_0,
0202                                                            std::vector<float> &input_model_state_memory_d_1_1) {
0203   /**************** Preprocessing ******************/
0204   // Create input tensor (including size and value) from the loaded inputs
0205   // Compute the product of all input dimension
0206   // Assign memory for input tensor
0207   // inputTensors will be used by the Session Run for inference
0208
0209   input_values.clear();
0210   input_values.emplace_back(digiHcalMapTW);
0211   input_values.emplace_back(numEvents);
0212   input_values.emplace_back(adThr);
0213   input_values.emplace_back(input_model_state_memory_e_0_0);
0214   input_values.emplace_back(input_model_state_memory_e_0_1);
0215   input_values.emplace_back(input_model_state_memory_e_1_0);
0216   input_values.emplace_back(input_model_state_memory_e_1_1);
0217   input_values.emplace_back(input_model_state_memory_d_0_0);
0218   input_values.emplace_back(input_model_state_memory_d_0_1);
0219   input_values.emplace_back(input_model_state_memory_d_1_0);
0220   input_values.emplace_back(input_model_state_memory_d_1_1);
0221
0222   /**************** Inference ******************/
0223
0224   output_values = ort_mSession->run(input_names, input_values, input_shapes, output_names, batch_size);
0225
0226   return output_values;
0227 }
0228
0229 // AD method to be called by the CMS system
0230 std::vector<std::vector<float>> OnlineDQMDigiAD::Inference_CMSSW(
0231     const std::vector<std::vector<float>> &digiHcal2DHist_depth_1,
0232     const std::vector<std::vector<float>> &digiHcal2DHist_depth_2,
0233     const std::vector<std::vector<float>> &digiHcal2DHist_depth_3,
0234     const std::vector<std::vector<float>> &digiHcal2DHist_depth_4,
0235     const std::vector<std::vector<float>> &digiHcal2DHist_depth_5,
0236     const std::vector<std::vector<float>> &digiHcal2DHist_depth_6,
0237     const std::vector<std::vector<float>> &digiHcal2DHist_depth_7,
0238     const float LS_numEvents,
0239     const float flagDecisionThr)
0240
0241 {
0242   /**************** Prepare data ******************/
0243   // merging all 2d hist into one 3d depth[ieta[iphi]]
0244   std::vector<std::vector<std::vector<float>>> digiHcal2DHist_depth_all;
0245
0246   if (hcal_subsystem_name == "he") {
0247     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_1);
0248     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_2);
0249     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_3);
0250     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_4);
0251     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_5);
0252     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_6);
0253     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_7);
0254   }
0255
0256   else if (hcal_subsystem_name == "hb") {
0257     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_1);
0258     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_2);
0259     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_3);
0260     digiHcal2DHist_depth_all.push_back(digiHcal2DHist_depth_4);
0261   }
0262
0263   // convert the 3d depth[ieta[iphi]] vector into 1d and commbined
0264   std::vector<float> digiHcalMapTW = PrepareONNXDQMMapVectors(digiHcal2DHist_depth_all);
0265
0266   std::vector<float> adThr{flagDecisionThr};  // AD decision threshold, increase to reduce sensitivity
0267   std::vector<float> numEvents{LS_numEvents};
0268
0269   // call model inference
0270   /**************** Inference ******************/
0271   std::vector<std::vector<float>> output_tensors = Inference(digiHcalMapTW,
0272                                                              numEvents,
0273                                                              adThr,
0274                                                              input_model_state_memory_e_0_0,
0275                                                              input_model_state_memory_e_0_1,
0276                                                              input_model_state_memory_e_1_0,
0277                                                              input_model_state_memory_e_1_1,
0278                                                              input_model_state_memory_d_0_0,
0279                                                              input_model_state_memory_d_0_1,
0280                                                              input_model_state_memory_d_1_0,
0281                                                              input_model_state_memory_d_1_1);
0282
0283   // auto output_tensors = Inference(digiHcalMapTW, numEvents, adThr);
0284   //std::cout << "******* model inference is success *******" << std::endl;
0285
0286   /**************** Output post processing ******************/
0287   //  split outputs into ad output vectors and state_memory vectors
0288   std::string state_output_name_tag = "rnn_hidden";
0289   std::vector<std::vector<float>> ad_model_output_vectors, ad_model_state_vectors;
0290   for (size_t i = 0; i < output_tensors.size(); i++) {
0291     std::string output_names_startstr = output_names[i].substr(
0292         2, state_output_name_tag.length());  // Extract the same number of characters as str2 from mOutputNames
0293     if (output_names_startstr == state_output_name_tag) {
0294       ad_model_state_vectors.emplace_back(output_tensors[i]);
0295     } else {
0296       ad_model_output_vectors.emplace_back(output_tensors[i]);
0297     }
0298   }
0299
0300   if (ad_model_output_vectors.size() == num_state_vectors) {
0301     input_model_state_memory_e_0_0 = ad_model_state_vectors[0];
0302     input_model_state_memory_e_0_1 = ad_model_state_vectors[1];
0303     input_model_state_memory_e_1_0 = ad_model_state_vectors[2];
0304     input_model_state_memory_e_1_1 = ad_model_state_vectors[3];
0305     input_model_state_memory_d_0_0 = ad_model_state_vectors[4];
0306     input_model_state_memory_d_0_1 = ad_model_state_vectors[5];
0307     input_model_state_memory_d_1_0 = ad_model_state_vectors[6];
0308     input_model_state_memory_d_1_1 = ad_model_state_vectors[7];
0309   } else {
0310     std::cout << "Warning: the number of output state vectors does NOT equals to expected!. The states are set to  "
0311                  "default values."
0312               << std::endl;
0313     InitializeState();
0314   }
0315
0316   // # if onnx is returning serialized 1d vectors instead of vector of 3d vectors
0317   // aml score and flag are at index 5 and 7 of the vector ad_model_output_vectors: anomaly score: ad_model_output_vectors[5], anomaly flags: ad_model_output_vectors[7]
0318   /*
0319       selOutputIdx: index to select of the onnx output. e.g. 5 is the anomaly score and 7 is the anomaly flag (1 is with anomaly, 0 is healthy)
0320       std::vector<std::vector<std::vector<float>>> digiHcal3DHist_ANOMALY_FLAG = ONNXOutputToDQMHistMap(ad_model_output_vectors, 7)
0321       std::vector<std::vector<std::vector<float>>> digiHcal3DHist_ANOMALY_SCORE = ONNXOutputToDQMHistMap(ad_model_output_vectors, 5)
0322       */
0323
0324   // reduce counter for each ls call. due to onnx double datatype handling limitation that might cause precision error to propagate.
0325   if (--model_state_refresh_counter == 0)
0326     InitializeState();
0327
0328   return ad_model_output_vectors;
0329 }