如何在Microsoft.ML.OnnxRuntime中获取Tensor<float>的float*优化性能？

问题描述 #

在.NET Framework 4.8环境中使用Microsoft.ML.OnnxRuntime NuGet包推理YOLOv7模型，调用Session.Run得到Tensor<float>类型的输出结果。遍历1,867,320个float元素时，通过[]索引访问耗时近300ms（i7 10700 CPU），但相同操作在C++中仅需3ms。推测性能损耗来自Tensor的get/set方法，希望找到绕过索引直接访问内部缓冲区的方案，通过unsafe代码提升速度。

原C#代码 #

IDisposableReadOnlyCollection<DisposableNamedOnnxValue> output = Session.Run(input);
Stopwatch sw = new Stopwatch();
foreach (DisposableNamedOnnxValue res in output)
{
    sw.Start();
    Tensor<float> f_data = res.AsTensor<float>();
    for (int i = 0; i < 143640; i++) // 输出维度：{1,143640,13} 
         for (int j = 0; j < 13; j++)
               float a = f_data[0, i, j]; 
    sw.Stop(); // 耗时300-310ms
}

C++对比代码 #

void PostProcess(vector<Value>& ort_outputs) {
    const float* pdata = ort_outputs[0].GetTensorMutableData<float>();
    for (int i = 0; i < 143640 * 13; i++) {
        float v = pdata[i];
    } // 耗时2-4ms
}

解决方案 #

方案1：使用ReadOnlySpan（无需unsafe代码） #

Tensor<T>提供了TryGetReadOnlySpan方法，可以直接获取内存连续的只读区域，避免多维索引的计算和边界检查开销，性能接近C++实现：

IDisposableReadOnlyCollection<DisposableNamedOnnxValue> output = Session.Run(input);
Stopwatch sw = new Stopwatch();
foreach (DisposableNamedOnnxValue res in output)
{
    sw.Start();
    Tensor<float> f_data = res.AsTensor<float>();
    if (f_data.TryGetReadOnlySpan(out ReadOnlySpan<float> span))
    {
        // 一维遍历，对应原三维索引[0,i,j]的线性位置为 i*13 + j
        for (int k = 0; k < span.Length; k++)
        {
            float a = span[k];
            // 执行你的后处理逻辑
        }
    }
    sw.Stop();
}

方案2：使用unsafe代码直接获取指针 #

如果需要直接操作内存指针，可以结合MemoryMarshal和Unsafe类实现（需引用NuGet包System.Runtime.CompilerServices.Unsafe）：

IDisposableReadOnlyCollection<DisposableNamedOnnxValue> output = Session.Run(input);
Stopwatch sw = new Stopwatch();
foreach (DisposableNamedOnnxValue res in output)
{
    sw.Start();
    Tensor<float> f_data = res.AsTensor<float>();
    if (f_data.TryGetReadOnlySpan(out ReadOnlySpan<float> span))
    {
        unsafe
        {
            float* pdata = (float*)System.Runtime.CompilerServices.Unsafe.AsPointer(ref MemoryMarshal.GetReference(span));
            int totalCount = span.Length;
            for (int k = 0; k < totalCount; k++)
            {
                float a = pdata[k];
                // 执行你的后处理逻辑
            }
        }
    }
    sw.Stop();
}

性能提升原因 #

Tensor<T>的多维索引器每次访问都需要计算线性偏移量并做边界检查，这会带来大量额外开销；而通过ReadOnlySpan或直接获取指针后，遍历操作直接访问连续内存块，逻辑和C++一致，能大幅缩小性能差距。

内容的提问来源于stack exchange，提问作者BloodAndCat