c++实现argmax

c++实现argmax处理语义分割trt模型输出结果

一般分割模型输出的是nchw格式数据

void argmax(float* x, int* y, int rows, int cols, int chan)
{
    for(int i = 0,wh = rows * cols; i < rows; i++)
    {
        for(int j = 0; j < cols; j++)
        {
            int max = -10000000;
            int max_ind = -1;
            for(int k = 0; k < chan; k++)
            {
                if(x[k * wh + i * cols + j] > max)
                {
                    max = x[k * wh + i * cols + j];
                    max_ind = k;
                }
                
            }
            y[i * cols + j] = max_ind;
        }
    }
}