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;