Let onnxRuntime handles output memory

Author: demoulinv

src/aliceVision/segmentation/segmentation.cpp

@@ -67,11 +67,7 @@ bool Segmentation::initialize()
         _ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, _parameters.modelWeights.c_str(), ortSessionOptions);
     #endif
 
-        Ort::MemoryInfo memInfoCuda("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemType::OrtMemTypeDefault);
-        Ort::Allocator cudaAllocator(*_ortSession, memInfoCuda);
-
         _output.resize(_parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
-        _cudaInput = cudaAllocator.Alloc(_output.size() * sizeof(float));
 #endif
     }
     else
@@ -87,20 +83,6 @@ bool Segmentation::initialize()
     return true;
 }
 
-bool Segmentation::terminate()
-{
-    if (_parameters.useGpu)
-    {
-#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_ONNX_GPU)
-        Ort::MemoryInfo mem_info_cuda("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemType::OrtMemTypeDefault);
-        Ort::Allocator cudaAllocator(*_ortSession, mem_info_cuda);
-        cudaAllocator.Free(_cudaInput);
-#endif
-    }
-
-    return true;
-}
-
 bool Segmentation::processImage(image::Image<IndexT>& labels, const image::Image<image::RGBfColor>& source)
 {
     // Todo : handle orientation and small images smaller than model input
@@ -299,34 +281,23 @@ bool Segmentation::processTile(image::Image<ScoredLabel>& labels, const image::I
         return false;
     }
 
-    std::vector<float> output(_parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
-    int idx = 0;
-    for (int ch = 0; ch < _parameters.classes.size(); ch++)
-    {
-        for (int i = 0; i < _parameters.modelHeight; i++)
-        {
-            for (int j = 0; j < _parameters.modelWidth; j++)
-            {
-                const std::vector<int64_t> coords = {0, ch, i, j};
-                output[idx++] = outTensor[0].At<float>(coords);
-            }
-        }
-    }
-
     if (!labelsFromOutputTensor(labels, outTensor[0]))
     {
         return false;
     }
 
+    std::vector<float> output(_parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
+    auto *outTData = outTensor.front().GetTensorMutableData<float>();
+    output.assign(outTData, outTData + _parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
+
     return true;
 }
 
 bool Segmentation::processTileGPU(image::Image<ScoredLabel>& labels, const image::Image<image::RGBfColor>::Base& source)
 {
     ALICEVISION_LOG_TRACE("Process tile using gpu");
 #if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_CUDA)
-    Ort::MemoryInfo mem_info_cuda("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemType::OrtMemTypeDefault);
-    Ort::Allocator cudaAllocator(*_ortSession, mem_info_cuda);
+    Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);
 
     std::vector<const char*> inputNames{"input"};
     std::vector<const char*> outputNames{"output"};
@@ -335,41 +306,33 @@ bool Segmentation::processTileGPU(image::Image<ScoredLabel>& labels, const image
     std::vector<float> transformedInput;
     imageToPlanes(transformedInput, source);
 
-    cudaMemcpy(_cudaInput, transformedInput.data(), sizeof(float) * transformedInput.size(), cudaMemcpyHostToDevice);
-
-    Ort::Value inputTensors = Ort::Value::CreateTensor<float>(
-      mem_info_cuda, reinterpret_cast<float*>(_cudaInput), transformedInput.size(), inputDimensions.data(), inputDimensions.size());
+    std::vector<Ort::Value> inputTensors;
+    inputTensors.emplace_back(Ort::Value::CreateTensor<float>(memInfo,
+                                                              transformedInput.data(),
+                                                              transformedInput.size(),
+                                                              inputDimensions.data(),
+                                                              inputDimensions.size()));
 
     std::vector<Ort::Value> outTensor;
 
     try
     {
-        outTensor = _ortSession->Run(Ort::RunOptions{nullptr}, inputNames.data(), &inputTensors, 1, outputNames.data(), 1);
+        outTensor = _ortSession->Run(Ort::RunOptions{nullptr}, inputNames.data(), inputTensors.data(), 1, outputNames.data(), 1);
     }
     catch (const Ort::Exception& exception)
     {
         ALICEVISION_LOG_ERROR("ERROR running model inference: " << exception.what());
         return false;
     }
 
-    int idx = 0;
-    for (int ch = 0; ch < _parameters.classes.size(); ch++)
-    {
-        for (int i = 0; i < _parameters.modelHeight; i++)
-        {
-            for (int j = 0; j < _parameters.modelWidth; j++)
-            {
-                const std::vector<int64_t> coords = {0, ch, i, j};
-                _output[idx++] = outTensor[0].At<float>(coords);
-            }
-        }
-    }
-
     if (!labelsFromOutputTensor(labels, outTensor[0]))
     {
         return false;
     }
 
+    auto *outTData = outTensor.front().GetTensorMutableData<float>();
+    _output.assign(outTData, outTData + _parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
+
 #endif
 
     return true;

src/aliceVision/segmentation/segmentation.hpp

@@ -57,7 +57,7 @@ class Segmentation
         }
     }
 
-    virtual ~Segmentation() { terminate(); }
+    virtual ~Segmentation() {}
 
     /**
      * Process an input image to estimate segmentation
@@ -72,25 +72,13 @@ class Segmentation
      */
     bool initialize();
 
-    /**
-     * Onnx destruction code
-     */
-    bool terminate();
-
     /**
      * Assume the source image is the correct size
      * @param labels the output label image
      * @param source the input image to process
      */
     bool tiledProcess(image::Image<IndexT>& labels, const image::Image<image::RGBfColor>& source);
 
-    /**
-     * Transform model output to a label image
-     * @param labels the output labels image
-     * @param modeloutput the model output vector
-     */
-    bool labelsFromModelOutput(image::Image<ScoredLabel>& labels, const std::vector<float>& modelOutput);
-
     /**
      * Transform model output to a label image
      * @param labels the output labels image