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

目录

 

什么是Reduce算子

理解Reduce的过程

理解Reduce的算法

3维矩阵示例

CPU算法示例

总结

TBD

---

什么是Reduce算子

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

3维矩阵(N,H,W) 在H维度方向做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维矩阵作为例子把算法的思路梳理如下：

3维矩阵与内存中的布局

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

沿Axis=1方向计算ArgMax

求解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的内存地址

按照输入矩阵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 轴之前的所有维数 Dk 作为输入高维矩阵
• 每个输入高维矩阵可以计算出 Dk 个ArgMax结果，其中每个ArgMax结果是通过遍历axis=i轴上Di 个元素来计算

 

TBD

引入CUDA 并行计算后GPU版本的算法和它的计算过程会变成什么样子