【算法分析】Reduce类算子-ArgMax解析

最新推荐文章于 2025-03-04 17:40:40 发布

HaoBBNuanMM

最新推荐文章于 2025-03-04 17:40:40 发布

阅读量4.2k

点赞数 6

本文链接：https://blog.csdn.net/HaoBBNuanMM/article/details/110956558

版权

本文深入讲解了Reduce算子的工作原理及应用，特别是针对三维矩阵在指定维度上进行ArgMax计算的过程，并提供了详细的算法示例。

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

什么是Reduce算子

常见的如TensorFlow的tf.reduce_sum、MNN框架中的ArgMax/ArgMin，概况起来说就是对输入的多维张量（Tensor）数据，在某一维上执行特定的计算（比如sum、求Max/Min），从而达到降低维度的目的，下图先展示了一个典型的Reduce算子计算过程，本文围绕该过程对算法做详细说明

理解Reduce的过程

如上图所示，2维矩阵(n,m) 有x和y两个维度的方向，在x维度做reduce计算得到的结果是1维向量（n）即x维度reduce为1仅剩下y维度；在y维度做reduce计算得到的结果是1位向量（m）即y维度reduce为1仅剩下x维度；

通常在机器学习或者并行计算中采用Axis来指定输入矩阵的各个维度，把上图的例子泛化后描述：

输入的张量数据A 其中 rank(A)=n，（D0，D1，... Dn-1）其中Di表示Axis=i的维数；Reduce(A，Axis=i）= B 其中rank(B)=n-1， (D0, D1 ... , Di-1, Di+1, ... Dn-1）即Di被reduce为1
所以2维的矩阵在某个Axis上Reduce后成为1维矩阵，3维矩阵在某个Axis上Reduce后成为2维矩阵，丢失的那一维即是执行Reduce计算的对应Axis

理解ArgMax的算法

3维矩阵示例

Reduce的算法很多，这里以开源推理框架MNN 的ArgMax算子为例说明Reduce算子的一种ArgMax算子是如何在CPU上实现的（代码链接）

代码的本质是实现算法，所以在解释代码前先以3维矩阵作为例子把算法的思路梳理如下：

求ArgMax(A, Axis=1)可以得到B(N,W)，整个计算过程可以参考“什么是Reduce算子”章节的动画

求解B(N,W)矩阵中的元素bij：

参考上图，2维的坐标（i，j）如果不容易理解，其实将reduce为1的Axis=1的轴（H维度方向）加上，扩展成3维作为（i，1，j）就可以看到，bij 其实是对第i个输入矩阵，在第j列上沿着Axis=1的轴（H维度方向）求ArgMax做Reduce计算

按照输入矩阵A（N，H，W）在内存中按行优先存储的布局，索引到第i个输入矩阵的内存起始地址为： A + i*H*W ; 索引到第i个输入矩阵的第j列的内存起始地址为： A + i*H*W + j；索引到第i个输入矩阵的第j列的内存结束地址为：A + i*H*W + (H-1)*W + j ; 第i个输入矩阵的第j列的相邻两个元素之间的offset为 W

所以如上图所示，整个算法的核心思路就是3个for循环依次遍历所有输入矩阵，每个输入矩阵的所有列，每一列的所有行，最后计算出B的每个元素bij

CPU算法示例

泛化上面的3维矩阵实例到一般情况，算法的CPU版本实现（代码链接）和描述如下

输入的张量数据A 其中 rank(A)=n，（D0，D1，... Dn-1）求ArgMax(A, Axis=i)

根据“理解Reduce的过程”章节可以知道计算结果=B，其中rank(B)=n-1， (D0, D1 ... , Di-1, Di+1, ... Dn-1）Di被reduce为1

核心算法的代码很短（仅考虑输入是NHWC的格式），但要看明白需要花一番功夫，这里解析如下

//求解ArgMax(input, axis)，求输入张量在axis轴上做ArgMax计算的结果

        if (mMode == ARGMAX) {
            auto srcOrigin = input->host<float>();//获得输入A的起始地址
            auto dstOrigin = output->host<int>();//获得输出B的起始地址
            for (int i = 0; i < mNum; ++i) {//第1个for循环，遍历所有输入的高维矩阵
                auto iptr = srcOrigin + i * mDim * mKeyExtent;//计算A中第i个高维矩阵的起始地址
                auto optr = dstOrigin + i * mKeyExtent;//计算B中存放第i个高维矩阵ArgMax计算结果的起始地址

                for(int k = 0; k < mKeyExtent; ++k){//第2个for循环，遍历第i个高维矩阵在axis轴之后所有的维数
                    int index      = 0;
                    float maxValue = -FLT_MAX;
                    for (int j = 0; j < mDim; ++j) {//第3个for循环，遍历axis轴上的每个元素做ArgMax计算
                        auto val = iptr[k + j * mKeyExtent];
                        if (val > maxValue) {
                            maxValue = val;
                            index    = j;
                        }
                    }
                    optr[k] = index;//保存计算结果到B
                }
            }

该算法最难理解的部分是3个for循环的条件参数，说明如下：

        const int dimensions = input->dimensions();//获得A的rank
        for (int i = 0; i < mAxis; ++i) {
            mNum = mNum * input->length(i);//计算A的axis轴之前所有的总维数，即axis轴之前所有输入的高维矩阵的总个数（相当于3维矩阵实例中的N）
        }
        mDim = input->length(mAxis);//进行ArgMax计算的axis轴的维数（相当于3维矩阵实例中的H）
        for (int i = mAxis + 1; i < dimensions; ++i) {
            mKeyExtent = mKeyExtent * input->length(i);//计算A的axis轴之后的总维数，即对于第i个高维矩阵，生成多少个计算结果（相当于3维矩阵实例中的W）
        }

总结

张量数据A 其中 rank(A)=n，（D0，D1，... Dn-1）求 ArgMax(A, axis=i) 的算法可以概况为3点

求解结果B 其中rank(B)=n-1， (D0, D1 ... , Di-1, Di+1, ... Dn-1）Di被Reduce为1并忽略（如果是Reduce为topK个元素，那么Di维度降为topK）
遍历axis=i 轴之前的所有维数 $\prod_{k=0}^{i-1}$ Dk 作为输入高维矩阵
每个输入高维矩阵可以计算出 $\prod_{k=i+1}^{n-1}$ Dk 个ArgMax结果，其中每个ArgMax结果是通过遍历axis=i轴上Di 个元素来计算