Updating ML model class with safe copies, avoiding memory overwrites

Author: ChSonnabend

CODEOWNERS

@@ -78,7 +78,7 @@
 /PWGUD @alibuild @amatyja @rolavick
 /PWGJE @alibuild @nzardosh @fjonasALICE @jaimenorman @mhemmer-cern
 /Tools/PIDML @alibuild @saganatt
-/Tools/ML @alibuild @fcatalan92 @fmazzasc
+/Tools/ML @alibuild @fcatalan92 @fmazzasc @ChSonnabend
 /Tutorials/PWGCF @alibuild @jgrosseo @victor-gonzalez @zchochul
 /Tutorials/PWGDQ @alibuild @iarsene @mcoquet642 @XiaozhiBai @mguilbau
 /Tutorials/PWGEM @alibuild @mikesas @rbailhac @dsekihat @ivorobye @feisenhu

Common/Tools/PID/pidTPCModule.h

@@ -46,6 +46,7 @@
 #include <TRandom.h>
 #include <TString.h>
 
+#include <algorithm>
 #include <chrono>
 #include <cstddef>
 #include <cstdint>
@@ -510,6 +511,7 @@ class pidTPCModule
     float duration_network = 0;
 
     std::vector<float> track_properties(track_prop_size);
+    std::vector<float> output_network; // output buffer, allocation is reused for all mass hypotheses
     uint64_t counter_track_props = 0;
     int loop_counter = 0;
 
@@ -601,14 +603,13 @@ class pidTPCModule
       }
 
       auto start_network_eval = std::chrono::high_resolution_clock::now();
-      float* output_network = network.evalModel(track_properties);
+      network.evalModel(track_properties, output_network);
       auto stop_network_eval = std::chrono::high_resolution_clock::now();
       duration_network += std::chrono::duration<float, std::ratio<1, 1000000000>>(stop_network_eval - start_network_eval).count();
-      for (uint64_t k = 0; k < prediction_size; k += output_dimensions) {
-        for (int l = 0; l < output_dimensions; l++) {
-          network_prediction[k + l + prediction_size * loop_counter] = output_network[k + l];
-        }
+      if (output_network.size() != prediction_size) {
+        LOG(fatal) << "Network output size (" << output_network.size() << ") does not match the expected prediction size (" << prediction_size << ")";
       }
+      std::copy(output_network.begin(), output_network.end(), network_prediction.begin() + prediction_size * loop_counter);
 
       counter_track_props = 0;
       loop_counter += 1;

PWGDQ/Tasks/quarkoniaToHyperons.cxx

@@ -1773,7 +1773,7 @@ struct QuarkoniaToHyperons {
       float k0shortScore = -1;
       if (mlConfigurations.calculateK0ShortScores) {
         // evaluate machine-learning scores
-        float* k0shortProbability = mlCustomModelK0Short.evalModel(inputFeatures);
+        const std::vector<float> k0shortProbability = mlCustomModelK0Short.evalModel(inputFeatures);
         k0shortScore = k0shortProbability[1];
       } else {
         k0shortScore = v0.k0ShortBDTScore();
@@ -1788,7 +1788,7 @@ struct QuarkoniaToHyperons {
       float lambdaScore = -1;
       if (mlConfigurations.calculateLambdaScores) {
         // evaluate machine-learning scores
-        float* lambdaProbability = mlCustomModelLambda.evalModel(inputFeatures);
+        const std::vector<float> lambdaProbability = mlCustomModelLambda.evalModel(inputFeatures);
         lambdaScore = lambdaProbability[1];
       } else {
         lambdaScore = v0.lambdaBDTScore();
@@ -1803,7 +1803,7 @@ struct QuarkoniaToHyperons {
       float antiLambdaScore = -1;
       if (mlConfigurations.calculateAntiLambdaScores) {
         // evaluate machine-learning scores
-        float* antilambdaProbability = mlCustomModelAntiLambda.evalModel(inputFeatures);
+        const std::vector<float> antilambdaProbability = mlCustomModelAntiLambda.evalModel(inputFeatures);
         antiLambdaScore = antilambdaProbability[1];
       } else {
         antiLambdaScore = v0.antiLambdaBDTScore();

PWGHF/TableProducer/candidateSelectorLcPidMl.cxx

@@ -307,12 +307,12 @@ struct HfCandidateSelectorLcPidMl {
         std::vector<double> inputFeaturesD{trackParPos1.getPt(), trackPos1.dcaXY(), trackPos1.dcaZ(), trackParNeg.getPt(), trackNeg.dcaXY(), trackNeg.dcaZ(), trackParPos2.getPt(), trackPos2.dcaXY(), trackPos2.dcaZ()};
         float scores[3] = {-1.f, -1.f, -1.f};
         if (dataTypeML == 1) {
-          auto* scoresRaw = model.evalModel(inputFeaturesF);
+          const auto scoresRaw = model.evalModel(inputFeaturesF);
           for (int iScore = 0; iScore < 3; ++iScore) {
             scores[iScore] = scoresRaw[iScore];
           }
         } else if (dataTypeML == 11) {
-          auto* scoresRaw = model.evalModel(inputFeaturesD);
+          const auto scoresRaw = model.evalModel(inputFeaturesD);
           for (int iScore = 0; iScore < 3; ++iScore) {
             scores[iScore] = scoresRaw[iScore];
           }

PWGLF/TableProducer/Strangeness/lambdakzeromlselection.cxx

@@ -210,19 +210,19 @@ struct lambdakzeromlselection {
 
     // calculate classifier output
     if (PredictLambda) {
-      float* LambdaProbability = lambda_bdt.evalModel(inputFeatures);
+      const std::vector<float> LambdaProbability = lambda_bdt.evalModel(inputFeatures);
       lambdaMLSelections(LambdaProbability[1]);
     }
     if (PredictGamma) {
-      float* GammaProbability = gamma_bdt.evalModel(inputFeatures);
+      const std::vector<float> GammaProbability = gamma_bdt.evalModel(inputFeatures);
       gammaMLSelections(GammaProbability[1]);
     }
     if (PredictAntiLambda) {
-      float* AntiLambdaProbability = antilambda_bdt.evalModel(inputFeatures);
+      const std::vector<float> AntiLambdaProbability = antilambda_bdt.evalModel(inputFeatures);
       antiLambdaMLSelections(AntiLambdaProbability[1]);
     }
     if (PredictKZeroShort) {
-      float* KZeroShortProbability = kzeroshort_bdt.evalModel(inputFeatures);
+      const std::vector<float> KZeroShortProbability = kzeroshort_bdt.evalModel(inputFeatures);
       kzeroShortMLSelections(KZeroShortProbability[1]);
     }
   }

PWGLF/TableProducer/Strangeness/strangenessbuilder.cxx

@@ -1006,7 +1006,7 @@ struct StrangenessBuilder {
                                          AvgPA,                             // 6. Avg Pointing Angle
                                          static_cast<float>(v0zRanks[ic])}; // 7. V0 Vtx z Rank
 
-        float* BDTProbability = deduplication_bdt.evalModel(inputFeatures);
+        const std::vector<float> BDTProbability = deduplication_bdt.evalModel(inputFeatures);
 
         if (BDTProbability[1] > bestMLScore) {
           bestMLScore = BDTProbability[1];

PWGLF/Tasks/Strangeness/derivedlambdakzeroanalysis.cxx

@@ -1877,7 +1877,7 @@ struct derivedlambdakzeroanalysis {
       float k0shortScore = -1;
       if (mlConfigurations.calculateK0ShortScores) {
         // evaluate machine-learning scores
-        float* k0shortProbability = mlCustomModelK0Short.evalModel(inputFeatures);
+        const std::vector<float> k0shortProbability = mlCustomModelK0Short.evalModel(inputFeatures);
         k0shortScore = k0shortProbability[1];
       } else {
         k0shortScore = v0.k0ShortBDTScore();
@@ -1892,7 +1892,7 @@ struct derivedlambdakzeroanalysis {
       float lambdaScore = -1;
       if (mlConfigurations.calculateLambdaScores) {
         // evaluate machine-learning scores
-        float* lambdaProbability = mlCustomModelLambda.evalModel(inputFeatures);
+        const std::vector<float> lambdaProbability = mlCustomModelLambda.evalModel(inputFeatures);
         lambdaScore = lambdaProbability[1];
       } else {
         lambdaScore = v0.lambdaBDTScore();
@@ -1907,7 +1907,7 @@ struct derivedlambdakzeroanalysis {
       float antiLambdaScore = -1;
       if (mlConfigurations.calculateAntiLambdaScores) {
         // evaluate machine-learning scores
-        float* antilambdaProbability = mlCustomModelAntiLambda.evalModel(inputFeatures);
+        const std::vector<float> antilambdaProbability = mlCustomModelAntiLambda.evalModel(inputFeatures);
         antiLambdaScore = antilambdaProbability[1];
       } else {
         antiLambdaScore = v0.antiLambdaBDTScore();

Tools/ML/MlResponse.h

@@ -190,8 +190,16 @@ class MlResponse
       LOG(fatal) << "Number of input nodes in the model " << mPaths[nModel] << " is different from the number of input features to be tested (" << numInputNodes << " vs " << numInputFeatures << ")";
     }
 
-    TypeOutputScore* outputPtr = mModels[nModel].template evalModel<TypeOutputScore>(input);
-    return std::vector<TypeOutputScore>{outputPtr, outputPtr + mNClasses};
+    // evalModel returns an owning copy of the (last) output tensor of the model
+    std::vector<TypeOutputScore> output = mModels[nModel].template evalModel<TypeOutputScore>(input);
+    if (output.size() < mNClasses) {
+      LOG(fatal) << "Model " << mPaths[nModel] << " returned " << output.size() << " scores, but " << static_cast<int>(mNClasses) << " classes are expected. Please check your configurables.";
+    }
+    if (output.size() > mNClasses) {
+      // keep only the first mNClasses scores (e.g. single-candidate probabilities of a multi-output model)
+      output.resize(mNClasses);
+    }
+    return output;
   }
 
   /// ML selections

Tools/ML/model.cxx

@@ -26,9 +26,11 @@
 #include <onnxruntime_c_api.h>
 #include <onnxruntime_cxx_api.h>
 
+#include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
+#include <iterator>
 #include <memory>
 #include <sstream>
 #include <string>
@@ -42,6 +44,9 @@ namespace ml
 
 std::string OnnxModel::printShape(const std::vector<int64_t>& v)
 {
+  if (v.empty()) {
+    return "[]";
+  }
   std::stringstream ss("");
   for (std::size_t i = 0; i < v.size() - 1; i++)
     ss << v[i] << "x";
@@ -90,7 +95,12 @@ void OnnxModel::initModel(const std::string& localPath, const bool enableOptimiz
 
   mEnv = std::make_shared<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "onnx-model");
   mSession = std::make_shared<Ort::Session>(*mEnv, modelPath.c_str(), sessionOptions);
+  mMemInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
 
+  mInputNames.clear();
+  mInputShapes.clear();
+  mOutputNames.clear();
+  mOutputShapes.clear();
   Ort::AllocatorWithDefaultOptions const tmpAllocator;
   for (std::size_t i = 0; i < mSession->GetInputCount(); ++i) {
     mInputNames.push_back(mSession->GetInputNameAllocated(i, tmpAllocator).get());
@@ -122,6 +132,95 @@ void OnnxModel::initModel(const std::string& localPath, const bool enableOptimiz
   LOG(info) << "--- Model initialized! ---";
 }
 
+std::vector<int64_t> OnnxModel::inferInputShape(const std::size_t iinput, const int64_t size) const
+{
+  const std::vector<int64_t>& modelShape = mInputShapes[iinput];
+
+  // Rank-1 input: the whole vector is the tensor
+  if (modelShape.size() < 2) {
+    return {size};
+  }
+
+  // Product of all non-batch dimensions; dynamic dimensions (< 0) cannot be inferred
+  int64_t totalSize = 1;
+  bool hasDynamicDim = false;
+  for (std::size_t idim = 1; idim < modelShape.size(); idim++) {
+    if (modelShape[idim] < 0) {
+      hasDynamicDim = true;
+    } else {
+      totalSize *= modelShape[idim];
+    }
+  }
+
+  if (hasDynamicDim) {
+    if (modelShape.size() == 2) {
+      // [batch, features] with dynamic feature dimension: interpret the vector as a single sample
+      return {1, size};
+    }
+    LOG(fatal) << "Input " << iinput << " (" << mInputNames[iinput] << ") has dynamic non-batch dimensions (" << printShape(modelShape) << "), the tensor shape cannot be inferred from a flat vector. Please provide std::vector<Ort::Value> inputs instead.";
+  }
+
+  if (totalSize <= 0 || size % totalSize != 0) {
+    LOG(fatal) << "Size of the input vector (" << size << ") is not a multiple of the model input size (" << totalSize << ") for input " << iinput << " (" << mInputNames[iinput] << ", shape " << printShape(modelShape) << ")";
+  }
+
+  std::vector<int64_t> inputShape;
+  inputShape.reserve(modelShape.size());
+  inputShape.push_back(size / totalSize);
+  for (std::size_t idim = 1; idim < modelShape.size(); idim++) {
+    inputShape.push_back(modelShape[idim]);
+  }
+  return inputShape;
+}
+
+std::vector<Ort::Value> OnnxModel::evalModelRaw(std::vector<Ort::Value>& input)
+{
+  if (!mSession) {
+    LOG(fatal) << "OnnxModel::evalModel called before initModel()";
+  }
+  if (input.size() != mInputNames.size()) {
+    LOG(fatal) << "Number of input tensors (" << input.size() << ") does not agree with the number of model inputs (" << mInputNames.size() << ")";
+  }
+  for (std::size_t i = 0; i < input.size(); i++) {
+    LOG(debug) << "Input tensor " << i << " shape: " << printShape(input[i].GetTensorTypeAndShapeInfo().GetShape());
+  }
+
+  std::vector<const char*> inputNamesChar(mInputNames.size(), nullptr);
+  std::transform(std::begin(mInputNames), std::end(mInputNames), std::begin(inputNamesChar),
+                 [](const std::string& str) { return str.c_str(); });
+
+  std::vector<const char*> outputNamesChar(mOutputNames.size(), nullptr);
+  std::transform(std::begin(mOutputNames), std::end(mOutputNames), std::begin(outputNamesChar),
+                 [](const std::string& str) { return str.c_str(); });
+
+  std::vector<Ort::Value> outputTensors;
+  try {
+    const Ort::RunOptions runOptions;
+    outputTensors = mSession->Run(runOptions, inputNamesChar.data(), input.data(), input.size(), outputNamesChar.data(), outputNamesChar.size());
+  } catch (const Ort::Exception& exception) {
+    LOG(fatal) << "Error running model inference: " << exception.what();
+  }
+
+  LOG(debug) << "Number of output tensors: " << outputTensors.size();
+  if (outputTensors.size() != mOutputNames.size()) {
+    LOG(fatal) << "Number of output tensors: " << outputTensors.size() << " does not agree with the model specified size: " << mOutputNames.size();
+  }
+  for (std::size_t i = 0; i < outputTensors.size(); i++) {
+    const std::vector<int64_t> shape = outputTensors[i].GetTensorTypeAndShapeInfo().GetShape();
+    LOG(debug) << "Output tensor " << i << " shape: " << printShape(shape);
+    bool shapeOk = (shape.size() == mOutputShapes[i].size());
+    for (std::size_t idim = 0; shapeOk && idim < shape.size(); idim++) {
+      // dynamic dimensions of the model (< 0) can take any value
+      shapeOk = (mOutputShapes[i][idim] < 0) || (shape[idim] == mOutputShapes[i][idim]);
+    }
+    if (!shapeOk) {
+      LOG(fatal) << "Shape of output tensor " << i << " does not agree with model specification! Output: " << printShape(shape) << " model: " << printShape(mOutputShapes[i]);
+    }
+  }
+
+  return outputTensors;
+}
+
 void OnnxModel::setActiveThreads(const int threads)
 {
   activeThreads = threads;