如何在Microsoft.ML中优化分割网络输出的处理效率？

分割网络输出形状为1×N×K的张量（N为像素总数，K为类别数），每个元素对应像素位置的类别得分。当前用于获取每个像素最优类别的代码可正常运行，但耗时是网络推理的4倍，代码如下：

private int[,] GetClasses(List<DisposableNamedOnnxValue> output)
{            
    Tensor<float> outTensor = output.First().AsTensor<float>();
    
    int[,] classes = new int[frameWidth,frameHeight];

    for (int i = 0; i < frameWidth; ++i)
    {
        for (int j = 0; j < frameHeight; ++j)
        {
            int finalClass = 0;
            float finalClassScore = 0;
            for (int k = 0; k < nClasses; ++k)
            {
                float score = outTensor[0, i * frameHeight + j, k];
                if (score > finalClassScore)
                {
                    finalClassScore = score;
                    finalClass = k;
                }
            }
            classes[i, j] = finalClass;
        }
    }

    return classes;
}

请问在Microsoft.ML中是否有更高效的实现方式？

在Microsoft.ML中可以利用内置的张量操作和硬件加速大幅提升性能，主要有两种优化思路：

1. 使用ArgMax内置操作直接获取类别索引 #

Microsoft.ML的Tensor类原生支持ArgMax方法，能在指定维度上快速找出最大值的索引。该方法底层借助SIMD指令或硬件加速实现，比手动嵌套循环效率高得多：

using Microsoft.ML.OnnxRuntime.Tensors;
using Microsoft.ML.OnnxRuntime;

private int[,] GetClassesOptimized(List<DisposableNamedOnnxValue> output)
{
    Tensor<float> outTensor = output.First().AsTensor<float>();
    
    // 在第2维度（类别维度，索引从0开始）执行ArgMax，得到每个像素的最优类别索引
    Tensor<long> argMaxTensor = outTensor.ArgMax(dimensions: new[] { 2 });
    
    // 将一维ArgMax结果转换为frameWidth×frameHeight的二维数组
    int[,] classes = new int[frameWidth, frameHeight];
    for (int i = 0; i < frameWidth; i++)
    {
        for (int j = 0; j < frameHeight; j++)
        {
            int idx = i * frameHeight + j;
            classes[i, j] = (int)argMaxTensor[0, idx];
        }
    }
    
    return classes;
}

2. 优化内存访问模式 #

原代码中i * frameHeight + j的索引计算会导致张量访问不连续，降低缓存命中率。可以先转置张量，让同一像素的类别得分连续存储，再配合ArgMax进一步提升效率：

private int[,] GetClassesWithTranspose(List<DisposableNamedOnnxValue> output)
{
    Tensor<float> outTensor = output.First().AsTensor<float>();
    
    // 转置张量，将维度从(1, N, K)转为(1, K, N)，让同一像素的类别得分连续排列
    Tensor<float> transposedTensor = outTensor.Transpose(new[] { 0, 2, 1 });
    
    // 在第1维度（原类别维度）执行ArgMax
    Tensor<long> argMaxTensor = transposedTensor.ArgMax(dimensions: new[] { 1 });
    
    // 转换为二维数组
    int[,] classes = new int[frameWidth, frameHeight];
    for (int i = 0; i < frameWidth; i++)
    {
        for (int j = 0; j < frameHeight; j++)
        {
            int idx = i * frameHeight + j;
            classes[i, j] = (int)argMaxTensor[0, idx];
        }
    }
    
    return classes;
}

额外优化建议 #

• 避免在循环内重复计算索引值，提前计算并复用可减少不必要的运算开销
• 如果使用GPU推理，OnnxRuntime会自动将ArgMax操作调度到GPU执行，无需额外代码，能进一步缩短耗时

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