[Infrastructure,Tools,Common,Tutorial,PWGDQ,PWGHF,PWGLF] [ML] Fix for memory issue in ML headers (#17836)

Co-authored-by: ALICE Action Bot <alibuild@cern.ch>

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

@@ -49,6 +49,7 @@
 #include <TRandom.h>
 #include <TString.h>
 
+#include <algorithm>
 #include <array>
 #include <chrono>
 #include <cstddef>
@@ -570,11 +571,12 @@ class pidTPCModule
     std::unique_ptr<float[]> networkPrediction(new float[predictionSize * NParticleTypes]); // For each mass hypotheses
 
     const float nNclNormalization = response->GetNClNormalization();
-    float durationNetwork = 0;
+    float durationNetwork = 0.f;
 
     std::vector<float> trackProperties(trackPropSize);
+    std::vector<float> outputNetwork; // output buffer, allocation is reused for all mass hypotheses
     uint64_t counterTrackProps = 0;
-    int loopCounter = 0;
+    uint64_t loopCounter = 0;
 
     // To load the Hadronic rate once for each collision
     std::vector<float> hadronicRateForCollision(collisions.size(), 0.0f);
@@ -626,14 +628,13 @@ class pidTPCModule
       }
 
       const auto startNetworkEval = std::chrono::high_resolution_clock::now();
-      const float* const outputNetwork = network.evalModel(trackProperties);
+      network.evalModel(trackProperties, outputNetwork);
       const auto stopNetworkEval = std::chrono::high_resolution_clock::now();
       durationNetwork += std::chrono::duration<float, std::ratio<1, NanoToOne>>(stopNetworkEval - startNetworkEval).count();
-      for (uint64_t kPrediction = 0; kPrediction < predictionSize; kPrediction += outputDimensions) {
-        for (int lOutputDim = 0; lOutputDim < outputDimensions; ++lOutputDim) {
-          networkPrediction[kPrediction + lOutputDim + predictionSize * loopCounter] = outputNetwork[kPrediction + lOutputDim];
-        }
+      if (outputNetwork.size() != predictionSize) {
+        LOG(fatal) << "Network output size (" << outputNetwork.size() << ") does not match the expected prediction size (" << predictionSize << ")";
       }
+      std::copy(outputNetwork.begin(), outputNetwork.end(), networkPrediction.get() + predictionSize * loopCounter);
 
       counterTrackProps = 0;
       ++loopCounter;

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

@@ -1893,7 +1893,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();
@@ -1908,7 +1908,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();
@@ -1923,7 +1923,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;