相关文章推荐
霸气的麦片  ·  AVFoundation 命名空间 | ...·  8 月前    · 
不敢表白的小马驹  ·  SINE: ...·  1 年前    · 

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;