在Microsoft.ML中处理分割网络输出的最佳方法是什么？

Question

网络产生1xNxK张量，其中N是像素位置数，K是类数，每个值代表给定位置的一个类的分数。

当前检索每个位置最佳类关联的代码正在工作，但它非常慢，而且x4比网络运行本身花费更多的时间。

        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中是否有更好、更快的方法来做到这一点？

Answer 1

我使用的解决方案是在初始的keras模型中添加argmax层。Keras output single value through argmax。