Tensor memory management

[Umenokin]

Hi folks,
Thank you for the great crate.

I want to use the Yolo model in pretty much real-time scenarios, thus wishing to avoid allocating memory in the heap on every frame.

From the documentation, I learned:

• ndarray's &CowArray - will always be copied.
• ndarray's ArcArray<T, D> and Array<T, D> - potentially copied
• ndarray's ArrayView<'_, T, D> - will always be copied.
• Arc<Box<[T]>>, Box<[T]>, or Vec<T> - will never be copied.

With that in mind, I would assume the safest for me would be just using Vec. Something like that:

    let my_vector: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
    let dims = [1usize, 2, 3];
    let tensor = Tensor::from_array((dims, my_vector))?;
``

1. In reality, that might not even matter because, according to the documentation, depending on the execution provider, data can end up on the GPU, for example, so even if Tensor has a reference to an original array, would it be fair to assume that original array should not be mutated?

2. What is the right way to update Tensor, avoiding redundant memory allocation? 
My assumption would be:
 - If I need to change only a few numbers, I should use data_ptr, correct?
 - If I want to update the entire array, like in the case of Yolo, I assume I would want to work with Tensor's memory directly, and in case of GPU, update my own cpu copy and then copy it into Tensor. Is their standard mechanism or something I need to manage myself using memory_info to understand the best strategy?
 
The same is true for output data. Yolo model returns 56x8400 bounding boxes. In the case of the CPU provider, should I iterate the tensor's data directly with methods like data_ptr and, in the case of GPU, use extract_tensor_mut to copy everything to avoid iterating through each output's item via bus? 

Cheers!

[decahedron1]

1. In reality, that might not even matter because, according to the documentation, depending on the execution provider, data can end up on the GPU, for example, so even if Tensor has a reference to an original array, would it be fair to assume that original array should not be mutated?

When using Session::run with a session allocated on a non-CPU device, the data in every input is always copied to the device before inference. Updating tensor will update the data on the GPU at the next point Session::run is called.

The only exception to this is IoBinding - a bound input must be manually synchronized in order for the data to be copied.

2. What is the right way to update Tensor, avoiding redundant memory allocation?
   
   My assumption would be:
   - If I need to change only a few numbers, I should use data_ptr, correct?
   - If I want to update the entire array, like in the case of Yolo, I assume I would want to work with Tensor's memory directly, and in case of GPU, update my own cpu copy and then copy it into Tensor. Is their standard mechanism or something I need to manage myself using memory_info to understand the best strategy?
   
   The same is true for output data. Yolo model returns 56x8400 bounding boxes. In the case of the CPU provider, should I iterate the tensor's data directly with methods like data_ptr and, in the case of GPU, use extract_tensor_mut to copy everything to avoid iterating through each output's item via bus?

The correct way would be to use Tensor::extract_tensor_mut. You should absolutely not be using data_ptr unless you need to (i.e., if you were allocating tensors directly on GPU).

Note that:

• CUDA requires that you copy the entire buffer¹ to/from host/device memory. The CPU cannot view GPU memory unless the whole buffer is copied from the device.
• Unless you use IoBinding or OutputSelector, session outputs will always be copied to the CPU, no matter what device the session ran on.

With this in mind:

1. There is no need to have a separate CPU copy of the input tensor, since it's on the CPU anyway.
2. extract_tensor will work just fine for the output since it's on the CPU.

---

For best performance I would also recommend two things:

1. Allocate the input tensor using Tensor::new on the CUDA_PINNED allocation device. This makes CPU -> GPU copies significantly faster.
2. If the shape of the output is known ahead of time, use OutputSelector to pre-allocate the output tensor.

Footnotes

1. Technically it is possible to only copy small views of a buffer, but in practice the overhead of multiple htod_copy calls would negate any benefit gained over a single htod_copy on the whole buffer. ↩

[Umenokin]

@decahedron1 got it, thanks. What about input values? I found plenty of examples in the docs about extracting data from the "output" tensor but could not find any information about input values? From the doc "For ONNX Runtime, a value represents any type that can be given to/returned from a session or operator." Where Values comes in three types one of which is Tensor.

I have a function to create tensor:

pub fn create_input_tensor(shape: &Vec<i64>) -> Result<Tensor<f32>> {
    // Calculate the size from the shape
    let size = shape.iter().map(|&item| item as usize).product();
    let allocator = Allocator::default();
    allocator.alloc::<f32>(size).ok_or_else(|| anyhow!("Failed to allocate memory"))?;
    let tensor = Tensor::new(&allocator, shape.as_slice())?;
    Ok(tensor)
}

But I cannot find an example to pass it to the run function. For Yolo model I am trying something like "self.model.run(ort::inputs!["images" => self.input_tensor])"

In the source their is an example "session.run(ort::inputs![Tensor::from_string_array(Array1::from_vec(vec!["hello", "world"]))?]?);" My alternative "self.model.run(ort::inputs![self.input_tensor])?" does not compile and even if it would, it is not clear how model would map it to "image" input.

I am wondering if data will be copied anyway I might use "self.model.run(ort::input!s["image" => my_vector])?"
However, from the source, it seems this would use the default allocator, which is always CPU, and even though I am currently experimenting with CoreML, my plan is to use a CUDA-supporting device.

[decahedron1]

It seems you are forgetting the ? after ort::inputs![], as it actually produces an ort::Result<SessionInputs<'_>>. It should be self.model.run(ort::inputs!["images" => self.input_tensor]?)?.

ort::inputs![] is simply a convenience macro that allows

• supplying inputs in either array format (where the input name is determined by the order of inputs in the graph - this can be determined via Netron or Session.inputs), or in map format where the input name is explicitly specified, like ort::inputs!["image" => tensor];
• supplying each input as either an ndarray::Array or as a Value[Ref[Mut]].

If you store input_tensor on self, note that you'll probably want to forego the ort::inputs![] macro and instead pass an array of SessionInputValues directly, like self.model.run([self.input_tensor.view().into()]). This is because ort::inputs![] doesn't support views, which we need to pass to prevent self.input_tensor from being moved out of self.

Whether data is copied on Session::run depends on what is actually passed:

• If an ndarray::Array/ArrayView/Vec<T>/other non-Value type is passed, then a copy may be performed, depending on the behavior specified in Tensor::from_array.
• If a Value or ValueRef[Mut] is passed, the data is not copied.

In either case, if the session resides on a device other than the CPU, the data is inevitably copied an additional time to the session's device, so an Array input may be copied twice per run, whereas a Value type may only ever be copied at most once per run. CPU Values passed to CPU Sessions will never be copied.

[Umenokin]

@decahedron1 got it, thanks a lot for clarifications. I seem to have managed to make two versions that would compile which reflect what you are saying:

let model_inputs = vec![("images", self.input_tensor.view())];
let outputs = self.model.run(model_inputs)?;

And because I now ended up creating small vectors every frame, I created SessionInput like this:

let session_input: SessionInputs = [("images", self.input_tensor.view())].into();
let outputs = self.model.run(model_inputs)?;

I am still experimenting. By tracing the source I interpret the line let session_input: SessionInputs = vec![("images", self.input_tensor.view())].into(); as this:

1. self.input_tensor.view() - turn Tensor into SessionInputValue::View
2. let session_input: SessionInputs = [("images", self.input_tensor.view())].into(); - turn this array into SessionInputs::ValueMap of tuple where the first value is Cow("images") and the second one is a reference to Tensor via SessionInputValue::View

To confirm my assumptions unwrapped farther:

let session_unwrap: SessionInputs = SessionInputs::ValueMap(vec![(Cow::Owned("images".to_string()), SessionInputValue::View(self.input_tensor.view().into_dyn()))]);

The last one compiles and model return Ok though haven't checked what kind of garbage I end up generating 😅