Ncnn 实现 Onnx Argmax 算子

老邱生

已于 2023-02-24 11:36:22 修改

阅读量900

点赞数

CC 4.0 BY-SA版权

文章标签：人工智能计算机视觉 opencv

于 2023-02-06 17:03:17 首次发布

本文链接：https://blog.csdn.net/2301_76382363/article/details/128900550

本文介绍了ncnn不完全支持ONNX算子的情况，特别是argmax算子不在ncnn的内置支持之列。接着，文章详细讲解如何利用ncnn::Mat实现ONNX的argmax函数，特别指出当前实现仅适用于batch_size为1的情况，若需扩展到其他batch大小，需要额外增加循环处理。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

1、 ncnn 支持的 onnx 算子

ncnn 并不能支持所有的onnx算，从 onnx2ncnn.cpp 源码中找出了其支持的 onnx 算子：

Constant，Dropout，BatchNormalization，BiasGelu，Clip，Conv，ConvTranspose，EmbedLayerNormalization，Gemm，GroupNorm，GRU，InstanceNormalization，LayerNorm，LSTM，MatMul，MultiHeadAttention，Pad，PRelu，Reshape，Resize，RNNSkipLayerNormalization，Slice，Upsample，adaptive_avg_pool2d，adaptive_max_pool2d，AveragePool，MaxPool .............

太多了，这里就不意义列举了。可以看到不支持 argmax 算子，当然不支持的算子还有很多，这里只谈论一种，后续不断的增加。

2、实现argmax函数

这里使用 ncnn::Mat 实现了 onnx 的 argmax 函数，其中要求 batch_size 为1，如果需要支持不为1的情况请自行修改代码，也就是在通道的 for 循环外再增加一个 batch 的循环，因为这里固定了只读取第0个 batch 的数据，通过遍历对比，得到每一个通道的最大值记录在maxVal中。

// batch = 1
void OnnxArgmax2Ncnn(const ncnn::Mat &cnMat, std::vector<float> argMax_data_, std::vector<float> argMax_idx_)
{
    struct timeval start, end;
    float time_use;
    float tempMax;
    int tempIndex;
    std::cout << cnMat.c << cnMat.h <<