Adjusting for better compute performance

Author: ChSonnabend

Tools/ML/model.cxx

@@ -26,11 +26,9 @@
 #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>
@@ -97,6 +95,8 @@ void OnnxModel::initModel(const std::string& localPath, const bool enableOptimiz
   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();
@@ -114,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]);
@@ -173,7 +181,7 @@ std::vector<int64_t> OnnxModel::inferInputShape(const std::size_t iinput, const
   return inputShape;
 }
 
-std::vector<Ort::Value> OnnxModel::evalModelRaw(std::vector<Ort::Value>& input)
+void OnnxModel::checkInput(const std::vector<Ort::Value>& input) const
 {
   if (!mSession) {
     LOG(fatal) << "OnnxModel::evalModel called before initModel()";
@@ -184,19 +192,15 @@ std::vector<Ort::Value> OnnxModel::evalModelRaw(std::vector<Ort::Value>& input)
   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> OnnxModel::evalModelRaw(std::vector<Ort::Value>& input)
+{
+  checkInput(input);
   std::vector<Ort::Value> outputTensors;
   try {
-    const Ort::RunOptions runOptions;
-    outputTensors = mSession->Run(runOptions, inputNamesChar.data(), input.data(), input.size(), outputNamesChar.data(), outputNamesChar.size());
+    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();
   }
@@ -206,21 +210,44 @@ std::vector<Ort::Value> OnnxModel::evalModelRaw(std::vector<Ort::Value>& input)
     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]);
-    }
+    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;

Tools/ML/model.h

@@ -45,8 +45,8 @@ namespace ml
 ///    for as long as the returned vector is alive.
 ///  - evalModel(input): convenience wrapper returning a copy of the *last* model output as std::vector<T>.
 ///  - 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).
+///    resized as needed and its capacity is reused across calls, which avoids repeated allocations of the result buffer.
+///    Only the last model output is requested from ONNX Runtime.
 ///
 /// Inputs given as std::vector<T> 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).
@@ -87,8 +87,8 @@ class OnnxModel
   template <typename T>
   void evalModel(std::vector<Ort::Value>& input, std::vector<T>& output)
   {
-    const std::vector<Ort::Value> outputTensors = evalModelRaw(input);
-    copyLastOutput<T>(outputTensors, output);
+    const Ort::Value tensor = evalModelLast(input);
+    copyOutput<T>(tensor, output);
   }
 
   /// Run a single-input model on a flat vector of features (batches are inferred from the model input shape)
@@ -178,6 +178,11 @@ class OnnxModel
   std::vector<std::string> mOutputNames;
   std::vector<std::vector<int64_t>> mOutputShapes;
 
+  // Pointers into the name vectors above, populated only after all names are loaded.
+  // Moving the model transfers the backing vectors without invalidating these pointers.
+  std::vector<const char*> mInputNamesChar;
+  std::vector<const char*> mOutputNamesChar;
+
   // Environment settings
   std::string modelPath;
   int activeThreads = 0;
@@ -197,20 +202,26 @@ class OnnxModel
     tensors.emplace_back(Ort::Value::CreateTensor<T>(mMemInfo, data.data(), data.size(), inputShape.data(), inputShape.size()));
   }
 
-  /// Copy the content of the last output tensor into output (reusing its allocation)
+  /// Request only the last output, keeping its tensor alive until the caller has copied it
+  Ort::Value evalModelLast(std::vector<Ort::Value>& input);
+  void checkInput(const std::vector<Ort::Value>& input) const;
+  void checkOutput(const Ort::Value& tensor, std::size_t index) const;
+
+  /// Copy the content of an output tensor into output (reusing its allocation)
   template <typename T>
-  void copyLastOutput(const std::vector<Ort::Value>& outputTensors, std::vector<T>& output) const
+  void copyOutput(const Ort::Value& tensor, std::vector<T>& 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<T>::type) {
       LOG(fatal) << "Requested output type (ONNX type id " << static_cast<int>(Ort::TypeToTensorType<T>::type) << ") does not match the model output tensor type (ONNX type id " << static_cast<int>(info.GetElementType()) << ")";
     }
     const T* data = tensor.GetTensorData<T>();
-    output.assign(data, data + info.GetElementCount());
+    const auto size = info.GetElementCount();
+    if (size == 0) {
+      output.clear();
+    } else {
+      output.assign(data, data + size);
+    }
   }
 
   // Internal function for printing the shape of tensors