From 8075964781bd06fdd0e4bb9f293cba90369ddacd Mon Sep 17 00:00:00 2001 From: Christian Sonnabend Date: Tue, 8 Sep 2026 20:18:37 +0200 Subject: [PATCH 1/2] Updating ML model class with safe copies, avoiding memory overwrites --- CODEOWNERS | 2 +- Common/Tools/PID/pidTPCModule.h | 11 +- PWGDQ/Tasks/quarkoniaToHyperons.cxx | 6 +- .../candidateSelectorLcPidMl.cxx | 4 +- .../Strangeness/lambdakzeromlselection.cxx | 8 +- .../Strangeness/strangenessbuilder.cxx | 2 +- .../derivedlambdakzeroanalysis.cxx | 6 +- Tools/ML/MlResponse.h | 12 +- Tools/ML/model.cxx | 99 ++++++++++ Tools/ML/model.h | 176 +++++++++++------- Tutorials/ML/applyOnnxModel.cxx | 8 +- 11 files changed, 247 insertions(+), 87 deletions(-) diff --git a/CODEOWNERS b/CODEOWNERS index 55b281330f3..0cb77420b32 100644 --- a/CODEOWNERS +++ b/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 diff --git a/Common/Tools/PID/pidTPCModule.h b/Common/Tools/PID/pidTPCModule.h index de97ef65eba..13a5b6c5f60 100644 --- a/Common/Tools/PID/pidTPCModule.h +++ b/Common/Tools/PID/pidTPCModule.h @@ -46,6 +46,7 @@ #include #include +#include #include #include #include @@ -510,6 +511,7 @@ class pidTPCModule float duration_network = 0; std::vector track_properties(track_prop_size); + std::vector 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>(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; diff --git a/PWGDQ/Tasks/quarkoniaToHyperons.cxx b/PWGDQ/Tasks/quarkoniaToHyperons.cxx index 9c1fcb459f8..84606f59fac 100644 --- a/PWGDQ/Tasks/quarkoniaToHyperons.cxx +++ b/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 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 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 antilambdaProbability = mlCustomModelAntiLambda.evalModel(inputFeatures); antiLambdaScore = antilambdaProbability[1]; } else { antiLambdaScore = v0.antiLambdaBDTScore(); diff --git a/PWGHF/TableProducer/candidateSelectorLcPidMl.cxx b/PWGHF/TableProducer/candidateSelectorLcPidMl.cxx index f27e5edc3cd..efccd81c0fd 100644 --- a/PWGHF/TableProducer/candidateSelectorLcPidMl.cxx +++ b/PWGHF/TableProducer/candidateSelectorLcPidMl.cxx @@ -307,12 +307,12 @@ struct HfCandidateSelectorLcPidMl { std::vector 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]; } diff --git a/PWGLF/TableProducer/Strangeness/lambdakzeromlselection.cxx b/PWGLF/TableProducer/Strangeness/lambdakzeromlselection.cxx index b27cb24c68c..3fd13a7aa95 100644 --- a/PWGLF/TableProducer/Strangeness/lambdakzeromlselection.cxx +++ b/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 LambdaProbability = lambda_bdt.evalModel(inputFeatures); lambdaMLSelections(LambdaProbability[1]); } if (PredictGamma) { - float* GammaProbability = gamma_bdt.evalModel(inputFeatures); + const std::vector GammaProbability = gamma_bdt.evalModel(inputFeatures); gammaMLSelections(GammaProbability[1]); } if (PredictAntiLambda) { - float* AntiLambdaProbability = antilambda_bdt.evalModel(inputFeatures); + const std::vector AntiLambdaProbability = antilambda_bdt.evalModel(inputFeatures); antiLambdaMLSelections(AntiLambdaProbability[1]); } if (PredictKZeroShort) { - float* KZeroShortProbability = kzeroshort_bdt.evalModel(inputFeatures); + const std::vector KZeroShortProbability = kzeroshort_bdt.evalModel(inputFeatures); kzeroShortMLSelections(KZeroShortProbability[1]); } } diff --git a/PWGLF/TableProducer/Strangeness/strangenessbuilder.cxx b/PWGLF/TableProducer/Strangeness/strangenessbuilder.cxx index 14f37772f44..0d034c5b1a8 100644 --- a/PWGLF/TableProducer/Strangeness/strangenessbuilder.cxx +++ b/PWGLF/TableProducer/Strangeness/strangenessbuilder.cxx @@ -1006,7 +1006,7 @@ struct StrangenessBuilder { AvgPA, // 6. Avg Pointing Angle static_cast(v0zRanks[ic])}; // 7. V0 Vtx z Rank - float* BDTProbability = deduplication_bdt.evalModel(inputFeatures); + const std::vector BDTProbability = deduplication_bdt.evalModel(inputFeatures); if (BDTProbability[1] > bestMLScore) { bestMLScore = BDTProbability[1]; diff --git a/PWGLF/Tasks/Strangeness/derivedlambdakzeroanalysis.cxx b/PWGLF/Tasks/Strangeness/derivedlambdakzeroanalysis.cxx index 93659300c3c..e3cf9c8f609 100644 --- a/PWGLF/Tasks/Strangeness/derivedlambdakzeroanalysis.cxx +++ b/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 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 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 antilambdaProbability = mlCustomModelAntiLambda.evalModel(inputFeatures); antiLambdaScore = antilambdaProbability[1]; } else { antiLambdaScore = v0.antiLambdaBDTScore(); diff --git a/Tools/ML/MlResponse.h b/Tools/ML/MlResponse.h index cdc5a130714..367667d0af6 100644 --- a/Tools/ML/MlResponse.h +++ b/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(input); - return std::vector{outputPtr, outputPtr + mNClasses}; + // evalModel returns an owning copy of the (last) output tensor of the model + std::vector output = mModels[nModel].template evalModel(input); + if (output.size() < mNClasses) { + LOG(fatal) << "Model " << mPaths[nModel] << " returned " << output.size() << " scores, but " << static_cast(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 diff --git a/Tools/ML/model.cxx b/Tools/ML/model.cxx index 6d2098068aa..8d8bc74ac76 100644 --- a/Tools/ML/model.cxx +++ b/Tools/ML/model.cxx @@ -26,9 +26,11 @@ #include #include +#include #include #include #include +#include #include #include #include @@ -42,6 +44,9 @@ namespace ml std::string OnnxModel::printShape(const std::vector& 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_LOGGING_LEVEL_WARNING, "onnx-model"); mSession = std::make_shared(*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 OnnxModel::inferInputShape(const std::size_t iinput, const int64_t size) const +{ + const std::vector& 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 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 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 OnnxModel::evalModelRaw(std::vector& 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 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 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 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 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; diff --git a/Tools/ML/model.h b/Tools/ML/model.h index 3be08e72fa9..06682acc945 100644 --- a/Tools/ML/model.h +++ b/Tools/ML/model.h @@ -25,11 +25,8 @@ #include #include -#include -#include #include #include -#include #include #include #include @@ -40,90 +37,112 @@ namespace o2 namespace ml { +/// \brief Thin wrapper around an ONNX Runtime session (CPU only) +/// +/// Inference entry points: +/// - evalModelRaw(): runs the session and returns the owning output tensors (one Ort::Value per model output). +/// This is the zero-copy path: the caller owns the tensors and may read them via GetTensorData() +/// for as long as the returned vector is alive. +/// - evalModel(input): convenience wrapper returning a copy of the *last* model output as std::vector. +/// - evalModel(input, output): same as above but writes into a caller-provided vector. The vector is +/// resized as needed and its capacity is reused across calls, which avoids per-call allocations in +/// hot loops (e.g. batched inference). +/// +/// Inputs given as std::vector are wrapped in an Ort::Value without copying. The input vector must therefore +/// stay alive until the call returns (which is always the case for the synchronous calls provided here). class OnnxModel { public: OnnxModel() = default; ~OnnxModel() = default; + // The ONNX session options are not copyable, so the model is move-only + OnnxModel(const OnnxModel&) = delete; + OnnxModel& operator=(const OnnxModel&) = delete; + OnnxModel(OnnxModel&&) = default; + OnnxModel& operator=(OnnxModel&&) = default; // Inferencing void initModel(const std::string&, const bool = false, const int = 0, const uint64_t = 0, const uint64_t = 0); - // template methods -- best to define them in header + /// Run the model on already prepared input tensors + /// \param input one Ort::Value per model input + /// \return the output tensors of the model (owning). Access to the data via output[i].GetTensorData() + std::vector evalModelRaw(std::vector& input); + + /// Run the model and copy the last output tensor into a vector + /// \param input one Ort::Value per model input + /// \return flattened content of the last output tensor template - T* evalModel(std::vector& input) + std::vector evalModel(std::vector& input) { - LOG(debug) << "Input tensor shape: " << printShape(input[0].GetTensorTypeAndShapeInfo().GetShape()); - // assert(input[0].GetTensorTypeAndShapeInfo().GetShape() == getNumInputNodes()); --> Fails build in debug mode, TODO: assertion should be checked somehow - - try { - const Ort::RunOptions runOptions; - std::vector 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 outputNamesChar(mOutputNames.size(), nullptr); - std::transform(std::begin(mOutputNames), std::end(mOutputNames), std::begin(outputNamesChar), - [&](const std::string& str) { return str.c_str(); }); - auto outputTensors = mSession->Run(runOptions, inputNamesChar.data(), input.data(), input.size(), outputNamesChar.data(), outputNamesChar.size()); - 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++) { - LOG(debug) << "Output tensor shape: " << printShape(outputTensors[i].GetTensorTypeAndShapeInfo().GetShape()); - if ((outputTensors[i].GetTensorTypeAndShapeInfo().GetShape() != mOutputShapes[i]) && (mOutputShapes[i][0] != -1)) { - LOG(fatal) << "Shape of tensor " << i << " does not agree with model specification! Output: " << printShape(outputTensors[i].GetTensorTypeAndShapeInfo().GetShape()) << " model: " << printShape(mOutputShapes[i]); - } - } - T* outputValues = outputTensors.back().GetTensorMutableData(); - return outputValues; - } catch (const Ort::Exception& exception) { - LOG(error) << "Error running model inference: " << exception.what(); - } - return nullptr; + std::vector output; + evalModel(input, output); + return output; } + /// Run the model and copy the last output tensor into the provided vector (allocation is reused between calls) + /// \param input one Ort::Value per model input + /// \param output vector which is filled with the flattened content of the last output tensor template - T* evalModel(std::vector& input) + void evalModel(std::vector& input, std::vector& output) { - const int64_t size = input.size(); - assert(size % mInputShapes[0][1] == 0); - std::vector inputShape{size / mInputShapes[0][1], mInputShapes[0][1]}; - std::vector inputTensors; - Ort::MemoryInfo memInfo = - Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); - inputTensors.emplace_back(Ort::Value::CreateTensor(memInfo, input.data(), size, inputShape.data(), inputShape.size())); - LOG(debug) << "Input shape calculated from vector: " << printShape(inputShape); - return evalModel(inputTensors); + const std::vector outputTensors = evalModelRaw(input); + copyLastOutput(outputTensors, output); + } + + /// Run a single-input model on a flat vector of features (batches are inferred from the model input shape) + /// \param input flattened input features, size must be a multiple of the number of input nodes + /// \return flattened content of the last output tensor + template + std::vector evalModel(std::vector& input) + { + std::vector output; + evalModel(input, output); + return output; } - // For 2D inputs + /// Run a single-input model on a flat vector of features (batches are inferred from the model input shape) + /// \param input flattened input features, size must be a multiple of the number of input nodes + /// \param output vector which is filled with the flattened content of the last output tensor template - T* evalModel(std::vector>& input) + void evalModel(std::vector& input, std::vector& output) { + if (mInputShapes.size() != 1) { + LOG(fatal) << "Model has " << mInputShapes.size() << " inputs but a single input vector was provided. Use the std::vector> or std::vector overload."; + } std::vector inputTensors; + inputTensors.reserve(1); + addInputTensor(inputTensors, input, 0); + evalModel(inputTensors, output); + } - Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); + /// Run a multi-input model: one flat vector of features per model input + /// \param input one flattened vector per model input + /// \return flattened content of the last output tensor + template + std::vector evalModel(std::vector>& input) + { + std::vector output; + evalModel(input, output); + return output; + } + /// Run a multi-input model: one flat vector of features per model input + /// \param input one flattened vector per model input + /// \param output vector which is filled with the flattened content of the last output tensor + template + void evalModel(std::vector>& input, std::vector& output) + { + if (input.size() != mInputShapes.size()) { + LOG(fatal) << "Model has " << mInputShapes.size() << " inputs but " << input.size() << " input vectors were provided."; + } + std::vector inputTensors; + inputTensors.reserve(input.size()); for (std::size_t iinput = 0; iinput < input.size(); iinput++) { - [[maybe_unused]] int totalSize = 1; - int64_t size = input[iinput].size(); - for (std::size_t idim = 1; idim < mInputShapes[iinput].size(); idim++) { - totalSize *= mInputShapes[iinput][idim]; - } - assert(size % totalSize == 0); - - std::vector inputShape{static_cast(size / totalSize)}; - for (std::size_t idim = 1; idim < mInputShapes[iinput].size(); idim++) { - inputShape.push_back(mInputShapes[iinput][idim]); - } - - inputTensors.emplace_back(Ort::Value::CreateTensor(memInfo, input[iinput].data(), size, inputShape.data(), inputShape.size())); + addInputTensor(inputTensors, input[iinput], iinput); } - - return evalModel(inputTensors); + evalModel(inputTensors, output); } // Reset session @@ -141,6 +160,7 @@ class OnnxModel int getNumInputNodes() const { return mInputShapes[0][1]; } std::vector> getInputShapes() const { return mInputShapes; } int getNumOutputNodes() const { return mOutputShapes[0][1]; } + std::vector> getOutputShapes() const { return mOutputShapes; } uint64_t getValidityFrom() const { return validFrom; } uint64_t getValidityUntil() const { return validUntil; } void setActiveThreads(const int); @@ -150,6 +170,7 @@ class OnnxModel std::shared_ptr mEnv = nullptr; std::shared_ptr mSession = nullptr; Ort::SessionOptions sessionOptions; + Ort::MemoryInfo mMemInfo{nullptr}; // CPU memory info, created once in initModel // Input & Output specifications of the loaded network std::vector mInputNames; @@ -163,8 +184,37 @@ class OnnxModel uint64_t validFrom = 0; uint64_t validUntil = 0; + /// Derive the tensor shape for input number iinput from the number of provided values + /// The first dimension is treated as batch dimension; all other dimensions are taken from the model + std::vector inferInputShape(const std::size_t iinput, const int64_t size) const; + + /// Wrap a flat data vector (no copy) into an Ort::Value and append it to tensors + template + void addInputTensor(std::vector& tensors, std::vector& data, const std::size_t iinput) const + { + const std::vector inputShape = inferInputShape(iinput, static_cast(data.size())); + LOG(debug) << "Input shape calculated from vector: " << printShape(inputShape); + tensors.emplace_back(Ort::Value::CreateTensor(mMemInfo, data.data(), data.size(), inputShape.data(), inputShape.size())); + } + + /// Copy the content of the last output tensor into output (reusing its allocation) + template + void copyLastOutput(const std::vector& outputTensors, std::vector& output) const + { + if (outputTensors.empty()) { + LOG(fatal) << "Model returned no output tensors"; + } + const Ort::Value& tensor = outputTensors.back(); + const auto info = tensor.GetTensorTypeAndShapeInfo(); + if (info.GetElementType() != Ort::TypeToTensorType::type) { + LOG(fatal) << "Requested output type (ONNX type id " << static_cast(Ort::TypeToTensorType::type) << ") does not match the model output tensor type (ONNX type id " << static_cast(info.GetElementType()) << ")"; + } + const T* data = tensor.GetTensorData(); + output.assign(data, data + info.GetElementCount()); + } + // Internal function for printing the shape of tensors - std::string printShape(const std::vector&); + static std::string printShape(const std::vector&); bool checkHyperloop(const bool = true); }; diff --git a/Tutorials/ML/applyOnnxModel.cxx b/Tutorials/ML/applyOnnxModel.cxx index 16151511df0..560e3f204bd 100644 --- a/Tutorials/ML/applyOnnxModel.cxx +++ b/Tutorials/ML/applyOnnxModel.cxx @@ -24,6 +24,8 @@ #include #include +#include +#include #include #include @@ -51,11 +53,11 @@ struct applyModel { void run(ProcessingContext& pc) { - // Here we evaluate the model - float* modelOutput = network.evalModel(modelInput); + // Here we evaluate the model. The output tensor is copied into a std::vector owned by us + const std::vector modelOutput = network.evalModel(modelInput); // And now we print the output - for (int i = 0; i < 5; i++) { + for (std::size_t i = 0; i < std::min(modelInput.size(), modelOutput.size()); i++) { LOG(info) << "Input: " << modelInput[i] << ", Output: " << modelOutput[i]; } pc.services().get().endOfStream(); From a31fa892c80f52c81d1a7c51c75d8a7131c3a9f8 Mon Sep 17 00:00:00 2001 From: Christian Sonnabend Date: Tue, 8 Sep 2026 22:34:27 +0200 Subject: [PATCH 2/2] Fixing build issue --- Tools/ML/MlResponse.h | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/Tools/ML/MlResponse.h b/Tools/ML/MlResponse.h index 1dcb52cb6b4..413534f48cb 100644 --- a/Tools/ML/MlResponse.h +++ b/Tools/ML/MlResponse.h @@ -229,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(input); - if (outputPtr == nullptr) { - LOG(fatal) << "Batched model evaluation failed for model " << mPaths[nModel]; + std::vector output = mModels[nModel].template evalModel(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(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{outputPtr, outputPtr + nRows * mNClasses}; + return output; } /// ML selections