[Infrastructure,Tools,Common,Tutorial,PWGDQ,PWGHF,PWGLF] [WIP] [ML] Fix for memory issue in ML headers

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

@@ -1891,7 +1891,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();
@@ -1906,7 +1906,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();
@@ -1921,7 +1921,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;
   }
 
   /// Get vector with model predictions for a batch of candidates
@@ -221,11 +229,16 @@ class MlResponse
       LOG(fatal) << "Number of input nodes in the model " << mPaths[nModel] << " differs from features per row (" << numInputNodes << " vs " << featuresPerRow << ")";
     }
 
-    TypeOutputScore* outputPtr = mModels[nModel].template evalModel<TypeOutputScore>(input);
-    if (outputPtr == nullptr) {
-      LOG(fatal) << "Batched model evaluation failed for model " << mPaths[nModel];
+    std::vector<TypeOutputScore> output = mModels[nModel].template evalModel<TypeOutputScore>(input);
+    const std::size_t expectedOutputSize = nRows * mNClasses;
+    if (output.size() < expectedOutputSize) {
+      LOG(fatal) << "Model " << mPaths[nModel] << " returned " << output.size() << " scores, but " << expectedOutputSize << " scores are expected for " << nRows << " rows and " << static_cast<int>(mNClasses) << " classes. Please check your configurables.";
+    }
+    if (output.size() > expectedOutputSize) {
+      // keep only the first scores (e.g. batched probabilities of a multi-output model)
+      output.resize(expectedOutputSize);
     }
-    return std::vector<TypeOutputScore>{outputPtr, outputPtr + nRows * mNClasses};
+    return output;
   }
 
   /// ML selections

Tools/ML/model.cxx

@@ -42,6 +42,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 +93,14 @@ 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);
 
+  mInputNamesChar.clear();
+  mOutputNamesChar.clear();
+  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());
@@ -104,6 +114,14 @@ void OnnxModel::initModel(const std::string& localPath, const bool enableOptimiz
   for (std::size_t i = 0; i < mSession->GetOutputCount(); ++i) {
     mOutputShapes.emplace_back(mSession->GetOutputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape());
   }
+  mInputNamesChar.reserve(mInputNames.size());
+  for (const auto& name : mInputNames) {
+    mInputNamesChar.push_back(name.c_str());
+  }
+  mOutputNamesChar.reserve(mOutputNames.size());
+  for (const auto& name : mOutputNames) {
+    mOutputNamesChar.push_back(name.c_str());
+  }
   LOG(info) << "Input Nodes:";
   for (std::size_t i = 0; i < mInputNames.size(); i++) {
     LOG(info) << "\t" << mInputNames[i] << " : " << printShape(mInputShapes[i]);
@@ -122,6 +140,114 @@ 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;
+}
+
+void OnnxModel::checkInput(const std::vector<Ort::Value>& input) const
+{
+  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<Ort::Value> OnnxModel::evalModelRaw(std::vector<Ort::Value>& input)
+{
+  checkInput(input);
+  std::vector<Ort::Value> outputTensors;
+  try {
+    const Ort::RunOptions runOptions{nullptr};
+    outputTensors = mSession->Run(runOptions, mInputNamesChar.data(), input.data(), input.size(), mOutputNamesChar.data(), mOutputNamesChar.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++) {
+    checkOutput(outputTensors[i], i);
+  }
+
+  return outputTensors;
+}
+
+Ort::Value OnnxModel::evalModelLast(std::vector<Ort::Value>& input)
+{
+  checkInput(input);
+  if (mOutputNamesChar.empty()) {
+    LOG(fatal) << "Model has no outputs";
+  }
+  Ort::Value output{nullptr};
+  try {
+    // A null RunOptions uses the runtime defaults without allocating options per call.
+    const Ort::RunOptions runOptions{nullptr};
+    mSession->Run(runOptions, mInputNamesChar.data(), input.data(), input.size(), &mOutputNamesChar.back(), &output, 1);
+  } catch (const Ort::Exception& exception) {
+    LOG(fatal) << "Error running model inference: " << exception.what();
+  }
+  checkOutput(output, mOutputShapes.size() - 1);
+  return output;
+}
+
+void OnnxModel::checkOutput(const Ort::Value& tensor, const std::size_t index) const
+{
+  const std::vector<int64_t> shape = tensor.GetTensorTypeAndShapeInfo().GetShape();
+  LOG(debug) << "Output tensor " << index << " shape: " << printShape(shape);
+  bool shapeOk = (shape.size() == mOutputShapes[index].size());
+  for (std::size_t idim = 0; shapeOk && idim < shape.size(); idim++) {
+    // Dynamic dimensions of the model (< 0) can take any value.
+    shapeOk = (mOutputShapes[index][idim] < 0) || (shape[idim] == mOutputShapes[index][idim]);
+  }
+  if (!shapeOk) {
+    LOG(fatal) << "Shape of output tensor " << index << " does not agree with model specification! Output: " << printShape(shape) << " model: " << printShape(mOutputShapes[index]);
+  }
+}
+
 void OnnxModel::setActiveThreads(const int threads)
 {
   activeThreads = threads;