164 lines
5.6 KiB
ReStructuredText
164 lines
5.6 KiB
ReStructuredText
Argmax
|
||
=================
|
||
|
||
|
||
沿指定轴查找最大 `topk` 个值的索引。当 `topk=1` 时，该算子等价于 `ArgMax`。
|
||
|
||
.. math::
|
||
|
||
Y_i = \underset{k}{\operatorname{argmax}} (X_{slice_i})
|
||
|
||
其中 :math:`X_{slice_i}` 是输入张量中沿指定轴的一个切片，函数返回该切片中最大值的索引 :math:`k`。
|
||
|
||
输入：
|
||
- **input** - 输入数据地址。
|
||
- **params** - 其他参数打包成数组。
|
||
- **core_mask** - 核掩码。
|
||
|
||
输出：
|
||
- **output** - 存储索引的输出张量。
|
||
- **output_value** - 如果 `return_values` 为 true，则此处存储找到的值。
|
||
|
||
支持平台：
|
||
``FT78NE``
|
||
``MT7004``
|
||
|
||
.. note::
|
||
- FT78NE 支持fp32
|
||
- MT7004 支持fp16, fp32
|
||
|
||
**参数数组结构：**
|
||
|
||
.. code-block:: c
|
||
:linenos:
|
||
|
||
long long params[12];
|
||
params[0] = (long long)in_shape;输入张量的维度信息数组。
|
||
params[1] = (long long)in_strides;输入张量的步长信息数组。
|
||
params[2] = (long long)out_strides;输出张量的步长信息数组。
|
||
params[3] = (long long)arg_elements;用于存放候选值的临时工作空间地址。
|
||
params[4] = (long long)index;用于存放候选索引的临时工作空间地址。
|
||
params[5] = (long long)topk;需要查找的最大值的数量。
|
||
params[6] = (long long)out_value;是否返回数值的标志。若为非0，则 `output_value` 必须提供有效地址。
|
||
params[7] = (long long)in_shape_size;输入张量的维度数 (即 `in_shape` 数组的长度)。
|
||
params[8] = (long long)axis;执行查找操作的轴。
|
||
|
||
**共享存储版本:**
|
||
|
||
.. c:function:: void fp_arg_max_s(float *input, void *output, float *output_value, long long *params, int core_mask)
|
||
.. c:function:: void hp_arg_max_s(half *input, void *output, half *output_value, long long *params, int core_mask)
|
||
|
||
**C调用示例：**
|
||
|
||
.. code-block:: c
|
||
:linenos:
|
||
:emphasize-lines: 21
|
||
|
||
//FT78NE示例
|
||
#include <stdio.h>
|
||
#include <argmax.h>
|
||
int main(int argc, char* argv[]) {
|
||
float* input = (float*)0x81000000; //需要初始化
|
||
void* output = (void*)0x82000000; //不需要初始化
|
||
float* output_value = (float*)0x83000000; //可选
|
||
float* arg_elements = (float*)0x84000000;//不需要初始化
|
||
int* index = (int*)0x85000000; //不需要初始化
|
||
|
||
int core_mask = 0b1111;
|
||
|
||
int *in_strides = (int*)0x86000000;
|
||
int *out_strides = (int*)0x86000200;
|
||
|
||
int in_shape_size = 4;//最多只考虑4维
|
||
int in_shape[4] = {4, 8, 16, 8};
|
||
|
||
int axis = 1; //要操作的维度,不能超过3
|
||
int topk = 3; //不超过in_shape[axis]
|
||
|
||
int out_value = 0; //是否输出值，1表示输出值，0表示输出索引
|
||
float *outputfp32 = (float *)output;
|
||
int *outputint = (int *)output;
|
||
|
||
srand(time(0));
|
||
// 初始化测试数据，包含各种情况
|
||
int i, j;
|
||
//tensor 1 int32
|
||
int in_total_elements = in_shape[0] * in_shape[1] * in_shape[2] * in_shape[3];
|
||
for(i = 0; i < in_total_elements; i ++) {
|
||
input[i] = (float)(rand()%100);
|
||
}
|
||
|
||
long long params[9];
|
||
params[0] = (long long)in_shape;
|
||
params[1] = (long long)in_strides;
|
||
params[2] = (long long)out_strides;
|
||
params[3] = (long long)arg_elements;
|
||
params[4] = (long long)index;
|
||
params[5] = (long long)topk;
|
||
params[6] = (long long)out_value;
|
||
params[7] = (long long)in_shape_size;
|
||
params[8] = (long long)axis;
|
||
|
||
fp_arg_max_s(input, output, output_value, params, core_mask);
|
||
return 0;
|
||
}
|
||
|
||
|
||
**私有存储版本:**
|
||
|
||
.. c:function:: void fp_arg_max_p(float *input, void *output, float *output_value, long long *params)
|
||
.. c:function:: void hp_arg_max_p(half *input, void *output, half *output_value, long long *params)
|
||
|
||
|
||
**C调用示例：**
|
||
|
||
.. code-block:: c
|
||
:linenos:
|
||
:emphasize-lines: 20-21
|
||
|
||
//FT78NE示例
|
||
#include <stdio.h>
|
||
#include <argmax.h>
|
||
int main(int argc, char* argv[]) {
|
||
float* input = (float*)0x10010000; //需要初始化
|
||
void* output = (void*)0x10020000; //不需要初始化
|
||
float* output_value = (float*)0x10030000; //可选
|
||
float* arg_elements = (float*)0x10040000;//不需要初始化
|
||
int* index = (int*)0x10050000; //不需要初始化
|
||
|
||
int *in_strides = (int*)0x1004E000;
|
||
int *out_strides = (int*)0x1004E200;
|
||
|
||
int in_shape_size = 4;//最多只考虑4维
|
||
int in_shape[4] = {4, 8, 16, 8};
|
||
|
||
int axis = 1; //要操作的维度,不能超过3
|
||
int topk = 3; //不超过in_shape[axis]
|
||
|
||
int out_value = 0; //是否输出值，1表示输出值，0表示输出索引
|
||
float *outputfp32 = (float *)output;
|
||
int *outputint = (int *)output;
|
||
|
||
srand(time(0));
|
||
// 初始化测试数据，包含各种情况
|
||
int i, j;
|
||
//tensor 1 int32
|
||
int in_total_elements = in_shape[0] * in_shape[1] * in_shape[2] * in_shape[3];
|
||
for(i = 0; i < in_total_elements; i ++) {
|
||
input[i] = (float)(rand()%100);
|
||
}
|
||
|
||
long long params[9];
|
||
params[0] = (long long)in_shape;
|
||
params[1] = (long long)in_strides;
|
||
params[2] = (long long)out_strides;
|
||
params[3] = (long long)arg_elements;
|
||
params[4] = (long long)index;
|
||
params[5] = (long long)topk;
|
||
params[6] = (long long)out_value;
|
||
params[7] = (long long)in_shape_size;
|
||
params[8] = (long long)axis;
|
||
|
||
fp_arg_max_p(input, output, output_value, params);
|
||
return 0;
|
||
} |