5. RKNN进阶使用说明

5.1 数据排列格式

目前 RKNN 的数据排列格式主要有以下四种：NHWC、NCHW、NCIHWC2、UNDEFINE。

其中 NHWC 和 NCHW 的数据排布为深度学习常见数据排列方式，本章节不做额外说明，重点讲述 RKNPU 硬件专用的 NCIHWC2 数据格式的存储以及转换。 ![](images/RKNN_SDK_高级用户指南/RKNPU NC1HWC2数据排布与存储.png) 图5-1 RKNPU NC1HWC2数据排布与存储

如图5-1所示，数字0代表一笔数据，即一次存放C2个数据，其中C2是由平台决定的，不同硬件平台的C2的规则约束由表5-1所示，C1为C/C2的上取整值。NCIHWC2数据存放的顺序与图中数值增长的顺序一致，先存放0 - 15的数据，再存放16 - 31的数据。以RK3568平台为例，当feature为(1,13,4,4)的int8数据，对应的NCIHWC2为(1,2,4,4,8)，此时C2为8，C1为2，feature在内存中在16 - 31排放的数据中，对应的每个C2数据块只有前5个数据有效，剩下的3个数据是额外补的对齐数据。

表5-1 不同硬件平台对应的C2值

RK2118 RK3566/RK3568 RK3588/RK3576 RV1103B/RV1106B RV1103/RV1106 RK3562
int8 / 8 16 8 16 16
float16 4 4 8 / 8 8

接下来重点介绍 NCIHWC2 数据排列转 NCHW 和 NHWC 数据在内存中的变化过程。

以 feature (1, 13, 2, 2) RK3568 为例，数据在内存排布中的转换，根据前文的对齐要求可知 feature(1, 13, 2, 2) 对应的 NCIHWC2 为 (1, 2, 2, 2, 8)，NCIHWC2 的存储如下图所示，红色部分为额外对齐的无效数据。 图5-2 NC1HWC2数据排布展开

移除无效数据转成 NCHW 即 (1, 13, 2, 2) 数据，在内存中的排布如下： 图5-3 NCHW数据排布

移除无效数据转成 NHWC 即 (1, 2, 2, 13) 数据，在内存中的排布如下： 图5-4 NHWC数据排布

转换示例代码：

NCIHWC2 转 NCHW：以 int8 数据排列的 NCIHWC2 转成 int8 数据排列的 NCHW 如下所示：

/*
*src:  表示NCIHWC2输入tensor的地址
*dst:  表示NCHW输出tensor的地址
*dims: 表示NCIHWC2的shape信息
*channel:表示NCHW输入的C的值
*h:  表示NCHW的h的值
*w:  表示NCHW的w的值
*/
int NCIHWC2_to_NCHW(const int8_t* src, int8_t* dst, int* dims, int channel, int h, int w)
{
    int batch = dims[0];
    int C1  = dims[1];
    int C2  = dims[4];
    int hw_src = dims[2] * dims[3];
    int hw_dst = h * w;
    for (int i = 0; i < batch; i++) {
        src = src + i * C1 * hw_src * C2;
        dst = dst + i * channel * hw_dst;
        for (int c = 0; c < channel; ++c) {
            int plane = c / C2;
            const int8_t* src_c = plane * hw_src * C2 + src;
            int offset = c % C2;
            for (int cur_h = 0; cur_h < h; ++cur_h)
            for (int cur_w = 0; cur_w < w; ++cur_w) {
                int cur_hw = cur_h * w + cur_w;
                dst[c * hw_dst + cur_h * w + cur_w] = src_c[C2 * cur_hw + offset];
            }
        }
    }
    return 0;
}

NCIHWC2 转 NHWC：以 int8 数据排列的 NCIHWC2 转成 int8 数据排列的 NHWC 如下所示：

/*
*src:  表示NCIHWC2输入tensor的地址
*dst:  表示NHWC输出tensor的地址
*dims: 表示NCIHWC2的shape信息
*channel:表示NHWC输入的C的值
*h:  表示NHWC的h的值
*w:  表示NHWC的w的值
*/
int NCIHWC2_to_NHWC(const int8_t* src, int8_t* dst, int* dims, int channel, int h, int w)
{
    int batch = dims[0];
    int C1  = dims[1];
    int C2  = dims[4];
    int hw_src = dims[2] * dims[3];
    int hw_dst = h * w;
    for (int i = 0; i < batch; i++) {
        src = src + i * C1 * hw_src * C2;
        dst = dst + i * channel * hw_dst;
        for (int cur_h = 0; cur_h < h; ++cur_h) {
            for (int cur_w = 0; cur_w < w; ++cur_w) {
                int cur_hw = cur_h * dims[3] + cur_w; 
        for (int c = 0; c < channel; ++c) {
            int plane = c / C2;
            const auto* src_c = plane * hw_src * C2 + src;
            int offset = c % C2;
            dst[cur_h * w * channel + cur_w * channel + c] = src_c[C2 * cur_hw + offset];
        }
    }
}
return 0;
}

5.2 RKNN Runtime 零拷贝调用

5.2.1 零拷贝介绍

目前在 RK3562 / RK3566 / RK3568 / RK3576 / RK3588 上有两组 API 可以使用，分别是通用 API 接口和零拷贝流程的 API 接口，RV1103 系列/RV1106 系列支持零拷贝流程的 API 接口。零拷贝示例代码可以参考 rknpu2 工程下的 examples/rknn_zero_copy 目录，有关 rknpu2 的详细内容和代码库请访问链接：https://github.com/airockchip/rknn-toolkit2/tree/master/rknpu2。

在推理 RKNN 模型时，原始数据要经过输入处理、NPU 运行模型、输出处理三大流程。目前根据不同模型输入格式和量化方式，接口内部会存在通用 API 和零拷贝 API 两种处理流程，如图5-5和图5-6所示，两组 API 的主要区别在于，通用接口每次更新帧数据，需要将外部模块分配的数据拷贝到 NPU 运行时的输入内存，而零拷贝流程的接口会直接使用预先分配的内存（包括 NPU 运行时创建的或外部其他框架创建的，比如 DRM 框架），减少了内存拷贝的花销，性能更优，带宽更少。当用户输入数据只有虚拟地址时，只能使用通用 API 接口；当用户输入数据有物理地址或 fd 时，两组接口都可以使用。通用 API 和零拷贝 API 不能混合调用。

图5-5 通用API的数据处理流程

图5-6 通零拷贝API数据处理流程

  1. 通用 API

通用 API 的流程如图5-5所示。对于数据的归一化、量化、数据排布格式转换、反量化在 CPU 上运行（在符合零拷贝输入要求的情况下，归一化和量化会运行在 NPU 上，但输入数据仍需要用 CPU 拷贝到模型的输入 buffer 上 ），模型本身的推理在 NPU 上运行。

  1. 零拷贝 API

零拷贝 API 的流程如图5-6所示。优化了通用 API 的数据处理流程，零拷贝 API 归一化、量化和模型推理都会在 NPU 上运行，NPU 输出的数据排布格式和反量化过程在 CPU 或者 NPU 上运行。零拷贝 API 对于输入数据流程的处理效率会比通用 API 高。

零拷贝场景的使用条件如下表所示：

表5-2 零拷贝输入要求

输入维度 输入对齐要求
RV1103B / RV1106B / RK3566 / RK3568 RK3562 / RK3576 / RK3588 / RV1103 / RV1106
4维，通道数是1、3、4 宽8字节对齐 宽16字节对齐
非4维 总大小8字节对齐 总大小16字节对齐

5.2.2 C API 零拷贝整体流程

零拷贝 API 接口使用 rknn_tensor_memory 结构体，需要在推理前创建并设置该结构体，并在推理后读取该结构体中的内存信息。根据用户是否需要自行分配模型的模块内存（输入/输出/权重/中间结果）和内存表示方式（文件描述符/物理地址等）差异，有下列三种典型的零拷贝调用流程，如图5-7至图5-9所示，红色部分表示专为零拷贝加入的接口和数据结构，斜体表示接口调用之间传递的数据结构。

  • 输入/输出内存由运行时分配

![](images/RKNN_SDK_高级用户指南/零拷贝 API接口调用流程(输入/输出内部分配).png) 图5-7 零拷贝 API接口调用流程(输入/输出内部分配)

![](images/RKNN_SDK_高级用户指南/零拷贝 API接口调用流程(输入/输出外部分配).png) 图5-8 零拷贝 API接口调用流程(输入/输出外部分配)

如图5-8所示，输入/输出内存由外部分配调用的是 rknn_create_mem_from_fd()/rknn_create_mem_from_phys() 接口创建 rknn_tensor_memory 结构体，rknn_set_io_mem() 设置输入输出 rknn_tensor_memory 结构体。

flush_cache 表示用户需要调用与分配的内存类型关联的接口来刷新输出缓存。

  • 输入/输出/权重/中间结果内存由外部分配

![](images/RKNN_SDK_高级用户指南/零拷贝 API接口调用流程(输入/输出/权重/中间结果外部分配).png) 图5-9 零拷贝 API接口调用流程(输入/输出外部分配)

如图5-9所示，输入/输出/权重/中间结果内存由外部分配调用的是 rknn_create_mem_from_fd()/rknn_create_mem_from_phys() 接口创建 rknn_tensor_memory 结构体，rknn_set_io_mem() 设置输入输出 rknn_tensor_memory 结构体，rknn_set_weight_mem()/rknn_set_internal_mem() 设置权重/中间结果 rknn_tensor_memory 结构体。

5.2.3 C API 零拷贝的用法

以图5-7零拷贝 API 接口调用流程（输入/输出内部分配）为例，用法如下：

  • rknn_query()

输入：

用 RKNN_QUERY_NATIVE_INPUT_ATTR 查询相关的属性（注意，不是 RKNN_QUERY_INPUT_ATTR ）。当查询出来的 fmt（或者称为 layout ）不同时，需要提前处理的方式也不一样。该方式查询出来的是输入硬件效率最优的 layout 和 type。

rknn_query() 输入的情况如下:

a. 当 layout 为 RKNN_TENSOR_NCHW 时，这种情况一般输入是4维，并且数据类型为 bool 或者 int64，当传数据给 NPU 时，也需要按照 NCHW 格式排列给 NPU。

b. 当 layout 为 RKNN_TENSOR_NHWC 时，这种情况一般输入是4维，并且数据类型为 float32/float16/int8/uint8，同时，输入通道数是1、3、4。当传数据给 NPU 时，也需要按照 NHWC 格式排列给 NPU。需要注意的是当 pass_through=1 时，width 可能需要做 stride 对齐，具体取决于查询出来的 w_stride 的值。

c. 当 layout 为 RKNN_TENSOR_NCIHWC2 时，这种情况一般输入是4维，并且数据类型为 float16/int8，同时，输入通道数不是1、3、4。当 pass_through=0 时，输入数据按照 NHWC 格式排列，接口内部会进行 NHWC 到 NCIHWC2 的 CPU 转换；当 pass_through=1 时，输入数据按照 NCIHWC2 格式排列，用户外部需转换好。

d. 当 layout 为 RKNN_TENSOR_UNDEFINED 时，这种情况一般输入不是4维，当传数据给 NPU 时，需要按照 ONNX 模型输入格式传给 NPU。NPU 不做任何的 mean/std 处理以及 layout 转换。

如果用户需要的输入配置不同于查询接口获取的 rknn_tensor_attr 结构体，可以对 rknn_tensor_attr 结构体进行对应修改，目前支持的可修改的输入数据类型如表5-3所示。特别注意：如果查询的数据类型是 uint8，用户想传入 float32 类型，则 rknn_tensor_attr 结构体的 size 要修改成原 size 的四倍，同时其中的数据类型要修改成 RKNN_TENSOR_FLOAT32。用该方式修改后硬件效率就不是最优了，接口内部会调用 CPU 进行数据类型转换。

表5-3 输入可修改的输入数据类型表

rknn_query 得到的模型数据类型

用户接口修改的数据类型

bool

int8

float16

int16

int32

Int64

bool

Y

int8

Y

uint8

Y

Y

float32

Y

Y

float16

Y

int16

Y

int32

Y

Int64

Y

输出：

用 RKNN_QUERY_NATIVE_OUTPUT_ATTR 查询相关的属性（注意，不是 RKNN_QUERY_OUTPUT_ATTR ）。当查询出来的 fmt（或者称为 layout ）不同时，需要后处理的方式也不一样。该方式查询出来的是输出硬件效率最优的 layout 和 type。

当输出是4维且用户需要 NHWC layout 的四维输出时，可以用 RKNN_QUERY_NATIVE_NHWC_OUTPUT_ATTR 查询相关的属性。该方式可以直接获得 NHWC layout 的输出。

rknn_query() 输出的情况如下:

a. 当 layout 为 RKNN_TENSOR_NCIHWC2 时，这种情况一般输出是4维，并且数据类型为 float16/int8。当用户需要 NCHW layout 时，外部需进行 NCIHWC2 到 NCHW 的 layout 转换。

b. 当 layout 为 RKNN_TENSOR_UNDEFINED 时，这种情况一般输出非4维，并且数据类型为 float16/int8。用户外部无需进行 layout 转换。

c. 当 layout 为 RKNN_TENSOR_NCHW 时，这种情况一般输出是4维，并且数据类型为 float16/int8。用户外部无需进行 layout 转换。

d. 当 layout 为 RKNN_TENSOR_NHWC 时，这种情况一般输出是4维，并且数据类型为 float16/int8。这种情况一般是用户调用 RKNN_QUERY_NATIVE_NHWC_OUTPUT_ATTR 接口查询出来的 layout。

如果用户需要的输出配置不同于查询接口获取的 rknn_tensor_attr 结构体，可以对 rknn_tensor_attr 结构体进行对应修改，可修改的配置信息如表5-4，表5-5所示，特别注意：如果查询输出的数据类型是 int8，用户想获取成 float32 类型输出，则 rknn_tensor_attr 结构体的 size 要修改成原 size 的四倍，同时其中的数据类型要修改成 RKNN_TENSOR_FLOAT32。用该方式修改后硬件效率就不是最优了，接口内部会调用 CPU 进行数据类型转换。

表5-4 输出可修改的输入数据类型表

rknn_query 得到的模型数据类型

用户接口修改的数据类型

bool

int8

float16

int16

int32

Int64

bool

Y

int8

Y

uint8

float32

Y

Y

float16

Y

int16

Y

int32

Y

Int64

Y

表5-5 输出可修改的 layout 类型表

rknn_query 查询得到的模型 layout 类型
NCIHWC2 NCHW NHWC UNDEFINE
用户接口设置的 layout 类型：NCIHWC2 Y
NCHW Y Y
NHWC Y
UNDEFINE Y

RK3562/RK3566/RK3568/RK3576/RK3588 支持的零拷贝接口 NPU 输出配置如表5-6所示，RV1106/RV1106B/RV1103/RV1103B 支持零拷贝接口 NPU 输出配置如表5-7所示。

表5-6 RK3562/RK3566/RK3568/RK3576/RK3588 零拷贝接口 NPU 支持的输出配置

模型类型 输出数据类型 输出维度 可支持 output layout
int8 模型 int8/float16/float32 4维 NCHW/NCIHWC2/NHWC
非4维 UNDEFINE
float16 模型 float16/float32 4维 NCHW/NCIHWC2/NHWC
非4维 UNDEFINE

表5-7 RV1103/RV1103B/RV1106/RV1106B 零拷贝接口 NPU 支持的输出配置

模型类型 输出数据类型 输出维度 可支持 output layout
int8 模型 int8/float16 4维 NCHW/NCIHWC2/NHWC
非4维 UNDEFINE

  • rknn_create_mem

零拷贝 API 接口使用 rknn_tensor_memory 结构体，需要在推理前创建并设置该结构体，并在推理后读取该结构体中的内存信息。当无需对 RKNN_QUERY_NATIVE_INPUT_ATTR，RKNN_QUERY_NATIVE_OUTPUT_ATTR 查询出来的 layout 和 type 进行修改时，直接采用默认配置的 size_with_stride 创建内存大小。若修改了相应的 layout 和 type，则需按照相应的 size 创建内存大小（例如输出的数据类型是 int8，用户想获取成 float32 类型输出，size 要修改成原 size 的四倍）。

  • rknn_set_io_mem

rknn_set_io_mem() 用于设置包含模型输入/输出内存信息的 rknn_tensor_mem 结构体，和 rknn_init() 类似，只要在最开始调用一次，后面反复执行 rknn_run() 即可。

5.3 NPU 多核配置

RK3588 通过3核 NPU，RK3576 通过2核 NPU 提供更强的算力。本章节将详细介绍多核 NPU 的配置方法，以提高模型的推理效率。

注：多核运行适用于网络层计算量较大的网络，对小网络提升幅度较小，甚至可能因为单核多核的切换（该切换需 CPU 介入）而导致性能下降。

5.3.1 多核运行配置方法

如果使用 Python 作为应用程序开发语言，可以通过 RKNN-Toolkit2 或 RKNN-Toolkit Lite2 init_runtime() 接口中的 “core_mask” 参数设置模型运行的 NPU 核心。该参数的详细说明如下表：

表5-8 init_runtime接口core_mask参数说明

参数 详细说明
core_mask 该参数用于设置模型运行的 NPU 核心。可选值和相应说明如下：

NPU_CORE_AUTO: 自动调度模式，模型将以单核模式自动运行在当前空闲的 NPU 核上。
NPU_CORE_0: 模型运行在 NPU Core0 上。
NPU_CORE_1: 模型运行在 NPU Core1 上。
NPU_CORE_2: 模型运行在 NPU Core2 上。
NPU_CORE_0_1: 模型同时运行在 NPU Core0 和 NPU Core1 上。
NPU_CORE_0_1_2: 模型同时运行在 NPU Core0，Core1 和 Core2 上。
NPU_CORE_ALL: 根据平台自动配置 NPU 核心数量。

默认值为 NPU_CORE_AUTO。

注：在 RKNN-Toolkit Lite2 上设置该参数时，值的前面要加上 RKNNLite，例如 RKNNLite.NPU_CORE_AUTO；如果在 RKNN-Toolkit2 上设置该参数时，值的前面要加上 RKNN，例如 RKNN.NPU_CORE_AUTO。

RKNN-Toolkit2 设置 NPU 核心，参考代码如下：

## Python
......
# Init runtime environment
print('--> Init runtime environment')
ret = rknn.init_runtime(target='rk3588', core_mask=RKNN.NPU_CORE_0)
if ret != 0:
    print('Init runtime environment failed')
    exit(ret)

RKNN-Toolkit Lite2 设置 NPU 核心，参考代码如下：

## Python
......
# Init runtime environment
print('--> Init runtime environment')
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
if ret != 0:
    print('Init runtime environment failed')
    exit(ret)
print('done')

如果使用 C/C++ 作为应用程序开发语言，可以调用 rknn_set_core_mask() 接口设置模型运行的 NPU 核心。该接口 core_mask 参数的详细说明如下表：

表5-9 rknn_set_core_mask接口core_mask参数说明

参数 详细说明
core_mask 该参数用于设置模型运行的NPU核心。可选值和相应说明如下：
RKNN_NPU_CORE_AUTO: 自动调度模式，模型将以单核模式自动运行在当前空闲的NPU核上。
RKNN_NPU_CORE_0: 模型运行在NPU Core0上。
RKNN_NPU_CORE_1: 模型运行在NPU Core1上。
RKNN_NPU_CORE_2: 模型运行在NPU Core2上。
RKNN_NPU_CORE_0_1: 模型同时运行在NPU Core0和NPU Core1上。
RKNN_NPU_CORE_0_1_2: 模型同时运行在NPU Core0，Core1和Core2上。
RKNN_NPU_CORE_ALL: 根据平台自动配置NPU核心数量。

使用C/C++ API设置模型运行NPU核心，参考代码如下：

// C++
// rknn_context ctx;
rknn_core_mask core_mask = RKNN_NPU_CORE_0;
int ret = rknn_set_core_mask(ctx, core_mask);

5.3.2 查看多核运行效果

本章节将详细说明 RKNN 模型以多核模式运行时的效果。

如果使用 RKNN-Toolkit2 连接开发板进行模型推理，需要在调用 rknn.init_runtime() 接口时将 perf_debug 参数设置成 True，接着调用 rknn.eval_perf() 接口，即可打印每层的运行信息。参考代码如下：

##Python
# Init runtime environment
ret = rknn.init_runtime(target='rk3588', device_id='29d5dd97766a5c27', perf_debug=True)
if ret != 0:
    print('Init runtime environment failed')
    exit(ret)

# Eval performance
rknn.eval_perf()

如果是直接在板端进行模型推理，需要在运行应用前将 RKNN_LOG_LEVEL 设成4或以上，此时将打印模型每层的运行信息。设置方法如下：

## Python
# 使用RKNN-Toolkit Lite2提供的Python接口，只需在创建RKNNLite对象时将verbose设成True即可
rknnlite = RKNNLite(verbose=True)

# 使用C/C++接口，则需要在运行二进制程序前设置如下环境变量
export RKNN_LOG_LEVEL=4

以 lenet 模型为例，通过上述设置后，终端将打印类似如下日志（为方便展示，删除了 InputShape, OutputShape, DDR Cycles, NPU Cycles, Total Cycles, Time(us), MacUsage(%), Task Number, Lut Number, RW(kb), FullName 等字段）：

ID

OpType

DataType

Target

WorkLoad(0/1/2)-Improve

Therical

1

InputOperator

UINT8

CPU

100.0%/0.0%/0.0% - Up

0.0%

2

Conv

UINT8

NPU

50.0%/50.0%/0.0% - Up

50.0%

3

MaxPool

INT8

NPU

100.0%/0.0%/0.0% - Up

0.0%

4

Conv

INT8

NPU

50.0%/50.0%/0.0% - Up

50.0%

5

MaxPool

INT8

NPU

100.0%/0.0%/0.0% - Up

0.0%

6

ConvRelu

INT8

NPU

48.1%/51.9%/0.0% - Up

48.1%

7

Conv

INT8

NPU

100.0%/0.0%/0.0% - Up

0.0%

8

Softmax

INT8

CPU

0.0%/0.0%/0.0% - Up

0.0%

9

OutputOperator

FLOAT16

CPU

0.0%/0.0%/0.0% - Up

0.0%

Total Operator Elapsed Time(us): 591

Total Memory RW Amount(MB): 0

模型每层运行信息中的 “WorkLoad(0/1/2)-ImproveTherical” 一列只在多核 NPU 上会打印，记录了模型每一层的任务在 NPU 核心上是如何分配以及其理论性能提升情况。例如 “50.0%/50.0%/0.0% - Up:50.0%” 代表该层的计算量以 Core0 负责50%，Core1 负责50%进行分配，该层的性能相比单核运行，理论能提升50%。如果某一层的性能没有提升，例如 “100.0%/0.0%/0.0% - Up:0.0%”，可能存在以下几种情况：

  • 该层的负载太小，小于 NPU 多核任务分配的粒度，因此该层运行在单核上；

  • 该类算子在 NPU 驱动中未实现多核任务切分，待后续版本支持。现有已支持多核任务切分的算子有：Conv, DepthwiseConvolution, Add, Concat, Relu, Clip, Relu6, ThresholdedRelu, PRelu, LeakyRelu。

5.3.3 多核性能提升技巧

可以尝试如下方法，以得到较高的多核运行性能：

  • 将 CPU/DDR/NPU 频率定到最高

  • 将应用绑定至 CPU 大核

  • 将 NPU 中断绑定至应用所对应的 CPU 大核

不同固件对应的定频命令有所区别，请参考 8.1.1 章节。

以将应用绑定到 CPU4 大核心为例，上面提到的后两点可以参考如下脚本：

interrupts=$(cat /proc/interrupts | grep npu)
interrupts_array=($interrupts)

irq1=$(echo ${interrupts_array[0]} | awk -F: '{print $1}')
irq2=$(echo ${interrupts_array[14]} | awk -F: '{print $1}')
irq3=$(echo ${interrupts_array[28]} | awk -F: '{print $1}')

for irq in $irq1 $irq2 $irq3; do
    echo 4 > /proc/irq/$irq/smp_affinity_list
done

taskset 10 /rknn_benchmark lenet.rknn "" 10 3 ## CPU4对应的 taskset 掩码值为 0x10

上述脚本会执行如下操作：

  • 执行 cat /proc/interrupts | grep npu 命令并解析出三个中断号（去除冒号）

  • 使用循环将每个中断号的 smp_affinity_list 设置为4（CPU4对应的ID为4）

  • 最后执行 taskset 10 /rknn_benchmark lenet.rknn "" 10 3 命令，CPU4对应的 taskset 参数为10（有关 taskset 的具体用法，请参考：https://man7.org/linux/man-pages/man1/taskset.1.html ）

通过上述操作，NPU 中断以及应用程序“rknn_benchmark”都将在 CPU4 上运行，这样可以消除 NPU 中断处理的核心切换开销。

5.4 动态 Shape

5.4.1 动态 Shape 功能介绍

动态 shape 是指模型输入数据的形状在运行时可以改变。它可以帮助处理输入数据大小不固定的情况，增加模型的灵活性。在之前仅支持静态 shape 的 RKNN 模型情况下，如果用户需要使用多个输入 shape，传统的做法是生成多个 RKNN 模型，在模型部署时初始化多个上下文分别执行推理，而在引入动态 shape 后，用户可以只保留一份与静态 shape RKNN 模型大小接近的动态 shape RKNN 模型，并使用一个上下文进行推理，从而节省 Flash 占用和 DDR 占用，动态 shape 在图像处理和序列模型推理中具有重要的作用，它的典型应用场景包括：

  • 序列长度改变的模型，常见于 NLP 模型，例如 BERT, GPT

  • 空间维度变化的模型，例如分割和风格迁移

  • 带 Batch 模型，Batch 维度上变化

  • 可变输出数量的目标检测模型

5.4.2 RKNN SDK 版本和平台要求

  • RKNN-Toolkit2 版本 >=1.5.0

  • RKNPU Runtime 库 (librknnrt.so) 版本 >=1.5.0

  • RK3562/RK3566/RK3568/RK3576/RK3588/RK3588S 平台的 NPU 支持该功能

5.4.3 生成动态 Shape 的 RKNN 模型

本节介绍使用 RKNN-Toolkit2 的 Python 接口生成动态 shape 的 RKNN 模型的步骤：

  1. 确认模型支持动态 shape

如果模型文件本身不是动态 shape，RKNN-Toolkit2 支持扩展成动态 shape 的 RKNN 模型。首先，用户要确认模型本身不存在限制动态 shape 的算子或子图结构，例如，常量的形状无法改变，RKNN-Toolkit2 工具在转换过程会报错，如果遇到不支持动态 shape 扩展的情况，用户要根据报错信息，修改模型结构，重新训练模型以支持动态 shape。建议使用原始模型本身就是动态 shape 的模型。

  1. 设置需要使用的输入形状

由于 NPU 硬件特性，动态 shape RKNN 模型不支持输入形状任意改变，要求用户设置有限个输入形状。对于多输入的模型，每个输入的 shape 个数要相同。例如，在使用 RKNN-Toolkit2 转换 Caffe 模型时，Python 代码示例如下：

## Python
dynamic_input = [
    [[1,3,224,224]],  # set the first shape for all inputs
    [[1,3,192,192]],  # set the second shape for all inputs
    [[1,3,160,160]],  # set the third shape for all inputs
]

# Pre-process config
rknn.config(mean_values=[[103.94, 116.78, 123.68]], std_values=[[58.82, 58.82, 58.82]], quant_img_RGB2BGR=True, dynamic_input=dynamic_shapes)

上述接口配置会生成支持3个 shape 分别是 [1,3,224,224]、[1,3,192,192] 和 [1,3,160,160] 的动态 shape RKNN 模型。

dynamic_input 中的 shape 与原始模型框架的 layout 一致。例如，对于相同的 224x224 大小的 RGB 图片做分类，TensorFlow/TFLite 模型输入是 [1,224,224,3]，而 ONNX 模型输入是 [1,3,224,224]。

  1. 量化

在设置好输入 shape 后，如果要做量化，则需要设置量化矫正集数据。工具会读取用户设置的最大分辨率输入做量化（是所有输入尺寸之和的最大的一组 shape）。例如，模型有两个输入，一个输入 shape 分别是 [1,224] 和 [1,12]，另一个输入 shape 分别 [1,40] 和 [1,80]，第一组 shape 所有输入尺寸之和是 1*224 + 1*40 = 264，第二组 shape 所有输入尺寸之和是 1*12 + 1*80 = 92，第一组 shape 所有输入尺寸之和更大，因此使用两个输入分别以 [1,224] 和 [1,40] 的 shape 做量化。

  • 如果量化矫正集是 jpg/png 图片格式，用户可以使用不同的分辨率的图片做量化，因为工具会对图片使用 opencv 的 resize 方法缩放到最大分辨率后做量化。

  • 如果量化矫正集是 npy 格式，则用户必须使用最大分辨率输入的 shape。量化后，模型内所有 shape 在运行时使用同一套量化参数进行推理。

另外，输入的最大分辨率 shape 在调用 rknn.config 时也会打印出来，如下：

W config: The 'dynamic_input' function has been enabled, the MaxShape is dynamic_input[0] = [[1,224],[1,40]]!
The following functions are subject to the MaxShape:
1. The quantified dataset needs to be configured according to MaxShape
2. The eval_perf or eval_memory return the results of MaxShape

  1. 推理评估或精度分析 动态 shape RKNN 模型做推理或做精度分析时，用户必须提供第2步中设置的其中一组 shape 的输入。接口使用上与静态 shape RKNN 模型场景一致，此处不做赘述。

完整的创建动态 shape RKNN 模型示例，请参考 https://github.com/airockchip/rknn-toolkit2/tree/master/rknn-toolkit2/examples/functions/dynamic_shape 。

5.4.4 C API 部署

得到动态 shape RKNN 模型后，接着使用 RKNPU2 C API 进行部署。按照接口形式，分为通用 API 和零拷贝 API 部署流程。

5.4.4.1 通用 API

使用通用 API 部署动态 shape RKNN 模型的流程如下图所示：

图5-10 动态shape输入接口的通用API调用流程

加载动态 shape RKNN 模型后，可以在运行时动态修改输入的 shape。首先，通过 rknn_query() 可以查询 RKNN 模型支持的输入 shape 列表，每个输入支持的 shape 列表信息以 rknn_input_range 结构体形式返回，它包含了每个输入的名称、数据布局信息、shape 个数以及具体 shape。接着，通过调用 rknn_set_input_shapes() 接口，传入包含每个输入 shape 信息的 rknn_tensor_attr 数组指针可以设置当前推理使用的 shape。在设置输入 shape 后，可以再次调用 rknn_query() 查询当前设置成功后的输入和输出 shape。

最后，按照通用 API 流程完成推理。每次切换输入 shape 时，需要再设置一次新的 shape，准备新 shape 大小的数据并再次调用 rknn_inputs_set() 接口。如果推理前不需要切换输入 shape，无需重复调用 rknn_set_input_shapes() 接口。

1. 初始化

调用 rknn_init() 接口初始化动态 shape RKNN 模型，

对于动态 shape RKNN 模型，在初始化上下文时有如下限制：

  • 不支持权重共享功能（带 RKNN_FLAG_SHARE_WEIGHT_MEM 标志的初始化）。

  • 不支持上下文复用功能（具体说明见 rknn_dup_context 接口）。

2. 查询 RKNN 模型支持的输入 shape 组合

初始化成功后，通过 rknn_query() 可以查询到 RKNN 模型支持的输入 shape 列表，每个输入支持的 shape 列表信息以 rknn_input_range 结构体形式返回，它包含了每个输入的名称，layout 信息，支持的 shape 个数以及具体 shape。C 代码示例如下：

// 查询模型支持的输入 shape
rknn_input_range dyn_range[io_num.n_input];
memset(dyn_range, 0, io_num.n_input * sizeof(rknn_input_range));
for (uint32_t i = 0; i < io_num.n_input; i++)
{
    dyn_range[i].index = i;
    ret = rknn_query(ctx, RKNN_QUERY_INPUT_DYNAMIC_RANGE, &dyn_range[i], sizeof(rknn_input_range));
    if (ret != RKNN_SUCC)
    {
        fprintf(stderr, "rknn_query error! ret=%d\n", ret);
        return -1;
    }
    dump_input_dynamic_range(&dyn_range[i]);
}

注意：对于多输入的模型，所有输入的 shape 按顺序一一对应，例如，有两个输入、多种 shape 的 RKNN 模型，第一个输入的第一个 shape 与第二个输入的第一个 shape 组合有效，不存在交叉的 shape 组合。例如，模型有两个输入 A 和 B，A 的 shape 分别是 [1,224] 和 [1,112]，B 的 shape 分别 [1,40] 和 [1,80]，此时，只支持以下两组输入 shape 的情况：

  • A shape = [1,224], B shape = [1,40]

  • A shape = [1,112], B shape = [1,80]

3. 设置输入 shape

在首次设置输入数据或者输入数据 shape 发生改变时，需要调用 rknn_set_input_shapes() 接口动态修改输入 shape。加载动态 shape RKNN 模型后，可以在运行时动态修改输入的 shape。通过调用 rknn_set_input_shapes() 接口，传入所有输入的 rknn_tensor_attr 数组，每个 rknn_tensor_attr 中的 dims, n_dims 和 fmt 三个成员信息表示了当前推理所用的 shape。C 代码示例如下：

/**
dynamic inputs shape range:
index=0, name=data, shape_number=2, range=[[1, 224, 224, 3],[1, 112, 224, 3]], fmt = NHWC
**/
input_attrs[0].dims[0] = 1;
input_attrs[0].dims[1] = 224;
input_attrs[0].dims[2] = 224;
input_attrs[0].dims[3] = 3;
input_attrs[0].fmt = RKNN_TENSOR_NHWC;
ret = rknn_set_input_shapes(ctx, io_num.n_input, input_attrs);
if (ret < 0)
{
    fprintf(stderr, "rknn_set_input_shapes error! ret=%d\n", ret);
    return -1;
}

其中，io_num.n_input 是输入数量，input_attrs 是模型输入的 rknn_tensor_attr 结构体数组。

注：这里设置的 shape 必须包含在第2步查询到的 shape 列表中。

在设置输入 shape 后，可以再次调用 rknn_query 查询当前设置成功后的输入和输出 shape，C 代码示例如下：

// 获取当前次推理的输入和输出 shape
rknn_tensor_attr cur_input_attrs[io_num.n_input];
memset(cur_input_attrs, 0, io_num.n_input * sizeof(rknn_tensor_attr));
for (uint32_t i = 0; i < io_num.n_input; i++)
{
    cur_input_attrs[i].index = i;
    ret = rknn_query(ctx, RKNN_QUERY_CURRENT_INPUT_ATTR, &cur_input_attrs[i], sizeof(rknn_tensor_attr));
    if (ret < 0)
    {
        printf("rknn_init error! ret=%d\n", ret);
        return -1;
    }
    dump_tensor_attr(&cur_input_attrs[i]);
}
rknn_tensor_attr cur_output_attrs[io_num.n_output];
memset(cur_output_attrs, 0, io_num.n_output * sizeof(rknn_tensor_attr));
for (uint32_t i = 0; i < io_num.n_output; i++)
{
    cur_output_attrs[i].index = i;
    ret = rknn_query(ctx, RKNN_QUERY_CURRENT_OUTPUT_ATTR, &cur_output_attrs[i], sizeof(rknn_tensor_attr));
    if (ret != RKNN_SUCC)
    {
        printf("rknn_query fail! ret=%d\n", ret);
        return -1;
    }
    dump_tensor_attr(&cur_output_attrs[i]);
}

注意事项：

  • rknn_set_input_shapes 接口要求输入 tensor 的 shape 为4维时，fmt 使用 NHWC，非4维时使用 UNDEFINED。

  • 在 rknn_set_input_shapes 尚未调用前，使用带 RKNN_QUERY_CURRENT 前缀的命令查询的 shape 信息是无效的。

4. 推理

在设置好当前输入 shape 后，假设输入 Tensor 的 shape 信息保存在 cur_input_attrs 数组中，以通用 API 接口为例，C 代码示例如下：

// 设置输入信息
rknn_input inputs[io_num.n_input];
memset(inputs, 0, io_num.n_input * sizeof(rknn_input));
for (int i = 0; i < io_num.n_input; i++)
{
    int height = cur_input_attrs[i].fmt == RKNN_TENSOR_NHWC ? cur_input_attrs[i].dims[1] : cur_input_attrs[i].dims[2];
    int width = cur_input_attrs[i].fmt == RKNN_TENSOR_NHWC ? cur_input_attrs[i].dims[2] : cur_input_attrs[i].dims[3];
    cv::resize(imgs[i], imgs[i], cv::Size(width, height));
    inputs[i].index = i;
    inputs[i].pass_through = 0;
    inputs[i].type = RKNN_TENSOR_UINT8;
    inputs[i].fmt = RKNN_TENSOR_NHWC;
    inputs[i].buf = imgs[i].data;
    inputs[i].size = imgs[i].total() * imgs[i].channels();
}

// 将输入数据转换成正确的格式后，放到输入缓冲区
ret = rknn_inputs_set(ctx, io_num.n_input, inputs);
if (ret < 0)
{
    printf("rknn_input_set fail! ret=%d\n", ret);
    return -1;
}
// 进行推理
printf("Begin perf ...\n");
double total_time = 0;
for (int i = 0; i < loop_count; ++i)
{
    int64_t start_us = getCurrentTimeUs();
    ret = rknn_run(ctx, NULL);
    int64_t elapse_us = getCurrentTimeUs() - start_us;
    if (ret < 0)
    {
        printf("rknn run error %d\n", ret);
        return -1;
    }
    total_time += elapse_us / 1000.f;
    printf("%4d: Elapse Time = %.2fms, FPS = %.2f\n", i, elapse_us / 1000.f, 1000.f * 1000.f / elapse_us);
}
printf("Avg FPS = %.3f\n", loop_count * 1000.f / total_time);

// 获取输出结果
rknn_output outputs[io_num.n_output];
memset(outputs, 0, io_num.n_output * sizeof(rknn_output));
for (uint32_t i = 0; i < io_num.n_output; ++i)
{
    outputs[i].want_float = 1;
    outputs[i].index = i;
    outputs[i].is_prealloc = 0;
}

ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);
if (ret < 0)
{
    printf("rknn_outputs_get fail! ret=%d\n", ret);
    return ret;
}

// 释放输出缓冲区buffer
ret = rknn_outputs_release(ctx, io_num.n_output, outputs);

5.4.4.2 零拷贝 API

对于零拷贝 API 而言，初始化成功后，通过 rknn_query() 可以查询 RKNN 模型支持的输入 shape 列表，调用 rknn_create_mem() 接口分配的输入和输出内存。接着，通过调用 rknn_set_input_shapes() 接口，传入包含每个输入 shape 信息的 rknn_tensor_attr 数组指针可以设置当前推理使用的 shape。在设置输入 shape 后，可以再次调用 rknn_query() 查询设置成功后的输入和输出 shape。最后，调用 rknn_set_io_mem() 接口设置需要的输入输出内存。每次切换输入 shape 时，需要再设置一次新的 shape，准备新 shape 大小的数据并再次调用 rknn_set_io_mem() 接口，如果推理前不需要切换输入 shape，无需重复调用 rknn_set_input_shapes() 接口。典型用法流程如下图所示：

图5-11 动态shape输入接口的零拷贝API调用流程 初始化、查询 RKNN 模型支持的输入 shape 组合、设置输入 shape 使用与上述通用 API 相同，此处不做赘述。不同之处在于，在设置输入 shape 后，使用的接口不同。零拷贝推理 C 代码示例如下：

// 创建最大的输入 tensor 内存
rknn_tensor_mem *input_mems[io_num.n_input];
for (int i = 0; i < io_num.n_input; i++) {
    // default input type is int8 (normalize and quantize need compute in outside)
    // if set uint8, will fuse normalize and quantize to npu
    input_attrs[i].type = RKNN_TENSOR_UINT8;
    // default fmt is NHWC, npu only support NHWC in zero copy mode
    input_attrs[i].fmt = RKNN_TENSOR_NHWC;
    input_mems[i] = rknn_create_mem(ctx, input_attrs[i].size_with_stride);
    ...
}

// 创建最大的输出 tensor 内存
rknn_tensor_mem *output_mems[io_num.n_output];
for (uint32_t i = 0; i < io_num.n_output; i++) {
    // default output type is depend on model, this require float32 to compute top5
    // allocate float32 output tensor
    int output_size = output_attrs[i].size * sizeof(float);
    output_mems[i] = rknn_create_mem(ctx, output_size);
    ...
}

// 加载输入并设置模型输入 shape，每次切换输入 shape 要调用一次
for (int s = 0; s < shape_num; ++s) {
    for (int i = 0; i < io_num.n_input; ++i) {
        for (int j = 0; j < input_attrs[i].n_dims; ++j) {
            input_attrs[i].dims[j] = shape_range[i].dyn_range[s][j];
        }
    }
    ret = rknn_set_input_shapes(ctx, io_num.n_input, input_attrs);
    if (ret < 0) {
        fprintf(stderr, "rknn_set_input_shapes error! ret=%d\n", ret);
        return -1;
    }
    ...
}

// 获取当前次推理的输入和输出 shape
printf("current input tensors:\n");
rknn_tensor_attr cur_input_attrs[io_num.n_input];
memset(cur_input_attrs, 0, io_num.n_input * sizeof(rknn_tensor_attr));
for (uint32_t i = 0; i < io_num.n_input; i++) {
    cur_input_attrs[i].index = i;
    // query info
    ret = rknn_query(ctx, RKNN_QUERY_CURRENT_INPUT_ATTR, &cur_input_attrs[i], sizeof(rknn_tensor_attr));
    if (ret < 0) {
        printf("rknn_init error! ret=%d\n", ret);
        return -1;
    }
    dump_tensor_attr(&cur_input_attrs[i]);
    ...
}

printf("current output tensors:\n");
rknn_tensor_attr cur_output_attrs[io_num.n_output];
memset(cur_output_attrs, 0, io_num.n_output * sizeof(rknn_tensor_attr));
for (uint32_t i = 0; i < io_num.n_output; i++) {
    cur_output_attrs[i].index = i;
    // query info
    ret = rknn_query(ctx, RKNN_QUERY_CURRENT_OUTPUT_ATTR, &cur_output_attrs[i], sizeof(rknn_tensor_attr));
    if (ret != RKNN_SUCC) {
        printf("rknn_query fail! ret=%d\n", ret);
        return -1;
    }
    ...
    dump_tensor_attr(&cur_output_attrs[i]);
}

// 指定NPU核心数量，仅3588/3576支持
rknn_set_core_mask(ctx, (rknn_core_mask){core_mask});

// 设置输入信息
rknn_input inputs[io_num.n_input];
memset(inputs, 0, io_num.n_input * sizeof(rknn_input));
std::vector<cv::Mat> resize_imgs;
resize_imgs.resize(io_num.n_input);
for (int i = 0; i < io_num.n_input; i++) {
    int height = cur_input_attrs[i].fmt == RKNN_TENSOR_NHWC ? cur_input_attrs[i].dims[1] : cur_input_attrs[i].dims[2];
    int width = cur_input_attrs[i].fmt == RKNN_TENSOR_NHWC ? cur_input_attrs[i].dims[2] : cur_input_attrs[i].dims[3];
    int stride = cur_input_attrs[i].w_stride;
    cv::resize(imgs[i], resize_imgs[i], cv::Size(width, height));
    int input_size = resize_imgs[i].total() * resize_imgs[i].channels();
    // 拷贝外部数据到零拷贝输入缓冲区
    if (width == stride) {
        memcpy(input_mems[i]->virt_addr, resize_imgs[i].data, input_size);
    } else {
        int height = cur_input_attrs[i].dims[1];
        int channel = cur_input_attrs[i].dims[3];
        // copy from src to dst with stride
        uint8_t *src_ptr = resize_imgs[i].data;
        uint8_t *dst_ptr = (uint8_t *)input_mems[i]->virt_addr;
        // width-channel elements
        int src_wc_elems = width * channel;
        int dst_wc_elems = stride * channel;
        for (int b = 0; b < cur_input_attrs[i].dims[0]; b++) {
            for (int h = 0; h < height; h++) {
                memcpy(dst_ptr, src_ptr, src_wc_elems);
                src_ptr += src_wc_elems;
                dst_ptr += dst_wc_elems;
            }
        }
    }
    ...
}

// 更新输入零拷贝缓冲区内存
for (int i = 0; i < io_num.n_input; i++) {
    cur_input_attrs[i].type = RKNN_TENSOR_UINT8;
    ret = rknn_set_io_mem(ctx, input_mems[i], &cur_input_attrs[i]);
    if (ret < 0) {
        printf("rknn_set_io_mem fail! ret=%d\n", ret);
        return -1;
    }
    ...
}

// 更新输出零拷贝缓冲区内存
for (uint32_t i = 0; i < io_num.n_output; i++) {
    // default output type is depend on model, this require float32 to compute top5
    cur_output_attrs[i].type = RKNN_TENSOR_FLOAT32;
    cur_output_attrs[i].fmt = RKNN_TENSOR_NCHW;
    // set output memory and attribute
    ret = rknn_set_io_mem(ctx, output_mems[i], &cur_output_attrs[i]);
    if (ret < 0) {
        printf("rknn_set_io_mem fail! ret=%d\n", ret);
        return -1;
    }
    ...
}

// 推理
printf("Begin perf ...\n");
double total_time = 0;
for (int i = 0; i < loop_count; ++i) {
    int64_t start_us = getCurrentTimeUs();
    ret = rknn_run(ctx, NULL);
    int64_t elapse_us = getCurrentTimeUs() - start_us;
    if (ret < 0) {
        printf("rknn run error %d\n", ret);
        return -1;
    }
    ...
}
total_time += elapse_us / 1000.f;
printf("%4d: Elapse Time = %.2fms, FPS = %.2f\n", i, elapse_us / 1000.f, 1000.f * 1000.f / elapse_us);
printf("Avg FPS = %.3f\n", loop_count * 1000.f / total_time);

注意事项：

  1. rknn_set_io_mem() 接口在动态 shape 情况下，输入 buffer 的 shape 和大小说明：

    • 初始化完成后和调用 rknn_set_input_shapes() 接口前，rknn_query() 接口使用 RKNN_QUERY_INPUT_ATTR 和 RKNN_QUERY_OUTPUT_ATTR 查询输入和输出 Tensor 的 shape 通常是最大的，用户可以使用这两个命令获取的大小来分配输入和输出内存。若遇到多输入模型，部分输入的 shape 可能不是最大的，此时需要搜索支持的 shape 中最大的规格，并分配最大的输入和输出内存。

    • 如果输入是非4维度，使用 fmt=UNDEFINED，传递原始模型输入 shape 的 buffer，大小则根据输入 shape 和 type 计算得到。

    • 如果输入是4维度，支持使用 fmt=NHWC 或者 NCIHWC2，传递 NHWC 或者 NCIHWC2 shape 和对应 size 的 buffer（通过 rknn_query 查询相应字段获取 shape 和 size）。

    • rknn_query() 接口中，标志位为 RKNN_QUERY_CURRENT_INPUT_ATTR 和 RKNN_QUERY_CURRENT_OUTPUT_ATTR 时获取原始模型输入/输出的 shape，其格式为 NHWC 或者 UNDEFINED；标志位为 RKNN_QUERY_CURRENT_NATIVE_INPUT_ATTR 和 RKNN_QUERY_CURRENT_NATIVE_OUTPUT_ATTR 时获取 NPU 以最优性能读取数据时模型输入/输出的 shape，其格式为 NHWC 或者 NCIHWC2。

  2. rknn_set_io_mem() 接口中使用的 buffer 排列格式为 NHWC 时，rknn_tensor_attr 中的 shape 和 fmt 需按照 RKNN_QUERY_CURRENT_INPUT_ATTR 查询到的信息进行设置；如果使用的 buffer 排列格式为 NCIHWC2 时，需要按照 RKNN_QUERY_CURRENT_NATIVE_INPUT_ATTR 查询到的信息进行设置。

完整的动态 shape C API Demo 请参考 https://github.com/airockchip/rknn-toolkit2/tree/master/rknpu2/examples/rknn_dynamic_shape_input_demo 。

5.5 自定义算子

5.5.1 自定义算子介绍

RKNN SDK 提供了一种自定义算子的机制，它允许开发者在 RKNN 模型的推理阶段定义和执行自定义的算子。通过实现自定义算子，开发者可以扩展模型功能，并且针对特定硬件（CPU 或者 GPU）进行优化，以充分利用硬件资源并提高推理速度。同时，开发自定义算子需要深刻的理解深度学习计算原理和目标硬件平台的特性，以确保正确性和性能。目前只支持 ONNX 模型自定义算子。

RKNN 自定义算子主要包括两大步骤：

  • 使用 RKNN-Toolkit2 注册自定义算子并导出 RKNN 模型。

  • 编写自定义算子的 C 代码实现，通过 RKNN API 加载注册并执行。 整体流程如下图所示：

图5-12 注册自定义算子的完整流程

5.5.2 整体流程介绍

5.5.2.1 使用 RKNN-Toolkit2 注册自定义算子并导出 RKNN 模型

  1. 准备 ONNX 模型：按照 ONNX 模型标准规范，用户设计自定义算子的 op 类型、名字、op 属性、输入/输出数量，并将该算子插入到 ONNX 中的拓扑图位置，用户使用 ONNX 包提供的 api 设计和导出 ONNX 模型。

  2. 实现自定义算子 Python 类：类里主要包括 shape_infer() 和 compute() 两个函数接口，并调用 rknn.reg_custom_op() 注册该算子。

  3. 构建与导出模型：如果 rknn.build() 执行成功，可以进行推理，否则需要检查步骤 2 的 shape_infer() 代码实现。然后运行仿真，如果仿真结果正确，可以调用 rknn.export_rknn() 接口导出 RKNN 模型，否则需要检查步骤 2 的 compute() 代码。

5.5.2.2 编写自定义算子的 C 代码实现，通过 RKNN API 加载注册并执行

  1. 根据 rknn_custom_op.h 的 rknn_custom_op 类，编写自定义算子的 C 代码实现，编写完成后，填写 rknn_custom_op 类的信息。

  2. 调用 rknn_register_custom_ops() 注册 rknn_custom_op 类的信息。

  3. 参考通用 API 或零拷贝 API 的流程，正常构建、推理模型，可以开启模型详细日志和 Dump 功能确认自定义算子实现的正确性。

5.5.2.3 使用 RKNN-Toolkit2 连板推理或精度分析

若用户需要对包含自定义算子的模型做 Python 连板精度分析，需要将自定义算子的回调函数实现代码编译成 so 后，放在指定的路径，并重启 RKNN Server。具体参考 5.5.4.4 章节。

5.5.3 Python 端处理

目前只有 ONNX 模型支持自定义算子，支持用户添加非 ONNX 标准的算子。

添加非 ONNX 标准的自定义算子用于新增一个不存在于 ONNX 算子列表内的新算子，该算子除了要满足 ONNX spec 规范以外，还要满足以下规则：

  • 算子的 op_type 不能与 ONNX 标准算子相同，推荐以 cst 字符开头。

  • 算子与其他算子必须要有连接关系，包含各个输入/输出的 shape，数据类型等。

  • 算子输入属性，支持 bool、int32、float32、int64 类型的单值或者数组。

  • 算子常量输入，支持 bool、int32、float32、int64 类型，该类型指未量化 ONNX 模型的数据类型。

因为非 ONNX 标准算子并不是 ONNX SPEC 内的标准算子，所以用户需要自行通过 ONNX 的 API 或其他框架的 API 来构建并导出一个包含非 ONNX 标准算子的 ONNX 模型。

这边为方便起见，以一个简单的修改 Softmax 的定义为示例，来构建一个包含非 ONNX 标准算子 cstSoftmax 的 ONNX 模型，修改方法如下：

import onnx

path="test_softmax.onnx"
model=onnx.load(path)

for node in model.graph.node:
    if node.op_type=="Softmax":
        node.op_type = "cstSoftmax"
        ... # 修改 cstSoftmax 的属性定义等
onnx.save(model,"./test_softmax_custom.onnx")

在构建完包含自定义算子的 ONNX 模型后，可使用 Netron 打开该 ONNX 模型，并检查该自定义算子是否符合自定义算子的规范，满足规范之后，就可以实现该自定义算子的算子类，具体实现如下（以自定义 Softmax 为例）：

import numpy as np
from rknn.api.custom_op import get_node_attr

class cstSoftmax:
    op_type = 'cstSoftmax'
    
    def shape_infer(self, node, in_shapes, in_dtypes):
        out_shapes = in_shapes.copy()
        out_dtypes = in_dtypes.copy()
        return out_shapes, out_dtypes
    
    def compute(self, node, inputs):
        x = inputs[0]
        axis = get_node_attr(node, 'axis')
        x_max = np.max(x, axis=axis, keepdims=True)
        tmp = np.exp(x - x_max)
        s = np.sum(tmp, axis=axis, keepdims=True)
        outputs = [tmp / s]
        return outputs

  • 自定义算子必须为一个 Python 类。

  • 自定义算子类必须包含一个名为 op_type 的字符串变量，与构建的 ONNX 中自定义算子类型名一致。

  • 自定义算子类必须包含成员函数 shape_infer(self, node, in_shapes, in_dtypes)，函数名、参数名都必须一致，否则报错。该函数用于自定义算子的 shape 推理，其中，node 为 ONNX 的算子节点对象，该对象里包含了自定义算子的属性和输入输出信息；in_shapes 为该算子所有输入的 shape 信息，格式为 [shape_0, shape_1, ...]，列表内的 shape 的类型为列表；in_dtypes 为该算子所有输入的 dtype 信息，格式为 [dtype_0, dtype_1, ...]，列表内的 dtype 的类型为 numpy 的 dtype 类型。另外该函数需要返回该算子所有输出的 shape 信息和 dtype 信息，格式与 in_shapes 和 in_dtypes 一致。

  • 自定义算子类必须包含成员函数 compute(self, node, inputs)，函数名和参数名都必须一致，否则报错。该函数用于自定义算子的推理。其中，node 为 ONNX 的算子节点对象，该对象里包含了自定义算子的属性和输入输出信息；inputs 为该算子的输入数据，格式为 [array_0, array_1, ...]，列表内的 array 的类型为 numpy 的 ndarray 类型。另外该函数需要返回该算子所有输出的数据，格式与 inputs 一致。

  • 如自定义算子含有自定义的属性，可通过 from rknn.api.custom_op import get_node_attr 来获取自定义算子的属性值。

在编写完自定义算子类后，可以通过 rknn.reg_custom_op() 进行算子类的注册，注册完后，就可以调用 rknn.build() 转换并生成 RKNN 模型。自定义算子类可以确保模型的转换和推理等功能的正常。具体实现如下（以自定义 Softmax 为例）：

from rknn.api import RKNN

# Create RKNN object
rknn = RKNN(verbose=True)

# Pre-process config
print('--> Config model')
rknn.config(mean_values=[[103.94, 116.78, 123.68]], std_values=[[58.82, 58.82, 58.82]], 
            quant_img_RGB2BGR=True, target_platform='rk3566')
print('done')

print('--> Register cstSoftmax op')
ret = rknn.reg_custom_op(cstSoftmax())
if ret != 0:
    print('Register cstSoftmax op failed!')
    exit(ret)
print('done')

print('--> Loading model')
ret = rknn.load_onnx(model='mobilenet_v2.onnx')
if ret != 0:
    print('Load model failed!')
    exit(ret)
print('done')

# Build model
print('--> Building model')
ret = rknn.build(do_quantization=True, dataset='./dataset.txt')
if ret != 0:
    print('Build model failed!')
    exit(ret)
    print('done')

rknn.reg_custom_op() 需要在 rknn.config() 和 rknn.load_xxx() 之间调用。

5.5.4 C API 部署

在得到带自定义算子的 RKNN 模型后，开始调用 C API 部署。首先，自定义算子的结构体和接口位于 rknn_custom_op.h 头文件，开发者程序需要包含该头文件。注册使用自定义算子的流程如下图所示： ![](images/RKNN_SDK_高级用户指南/注册自定义算子的C API调用流程.png) 图5-13 注册自定义算子的C API调用流程

5.5.4.1 初始化自定义算子结构体

创建 rknn 上下文后，开发者需要创建 rknn_custom_op 结构体，设置自定义算子信息。算子信息包含以下内容：

  • version：算子的版本号。

  • target：算子的执行后端，目前支持 CPU 和 GPU。

  • op_type：算子的类型，与 ONNX 模型中的类型字段相同。

  • cl_kernel_name：OpenCL 代码的 cl_kernel 函数名。注册 GPU 算子时必须配置。

  • cl_kernel_source：自定义算子的 .cl 文件全路径或者 OpenCL kernel 字符串。当 cl_source_size=0，它表示 .cl 文件全路径；当 cl_source_size>0，它表示 OpenCL kernel 代码字符串。注册 GPU 算子时必须配置。

  • cl_source_size：cl_kernel_source 的大小。大小等于 0 是特殊情况，它表示 cl_kernel_source 是路径。

  • cl_build_options：OpenCL kernel 编译选项，以字符串形式传入。注册 GPU 算子时必须配置。

  • init：可选，在 rknn_register_custom_ops 被调用一次。

  • prepare：可选，它是预处理回调函数，每次 rknn_run 都会执行 prepare 和 compute/compute_native 回调，执行顺序是 prepare 在前，compute/compute_native 在后。

  • compute：必须实现，算子运算回调函数，它的输入/输出都是 NCHW 的 float32 格式数据（ONNX 模型如果指定输入/输出为 int64 的数据类型，则 int64 格式数据）

  • compute_native：保留，请设置成 NULL 。

  • destroy：可选，rknn_destroy 中执行一次。

  • init/prepare/compute 回调函数参数定义规范如下：

    • rknn_custom_op_context* op_ctx：op 回调函数的上下文信息

    • rknn_custom_op_tensor* inputs：op 输入 tensor 数据和信息

    • uint32_t n_inputs：op 输入个数

    • rknn_custom_op_tensor* outputs：op 输出 tensor 数据和信息

    • uint32_t n_outputs：op 输出个数

  • destroy 回调函数仅 rknn_custom_op_context* op_ctx 一个参数。

rknn_custom_op_context
包含 target（执行后端设备）、GPU 上下文、自定义算子私有上下文以及 priv_data，其中 priv_data 由开发者自行管理（赋值，读写，销毁），GPU 上下文包含 cl_context、cl_command_queue、cl_kernel 指针，可以通过强制类型转换得到对应的 OpenCL 对象。
priv_data 是一个用户可选是否配置的指针，通常的用法是用户在 init() 回调函数内创建资源，并将 priv_data 指向该段内存地址，在 prepare()/compute() 回调函数中操作，最终在 destroy() 回调函数内销毁资源。

rknn_custom_op_tensor
表示输入/输出 tensor 的信息，包含 tensor 的名称、形状、大小、量化参数、虚拟基地址、fd、数据偏移等信息。
用户在回调 compute() 回调函数内无需创建该算子的输入和输出 tensor 内存。虚拟地址对应的数据在进入 compute() 回调函数时已经准备好。虚拟地址的计算公式是 Tensor 的有效地址 = 虚拟基地址 + 数据偏移，mem 成员的 virt_addr 表示虚拟基地址，mem 成员的 offset 表示数据偏移（以字节为单位）。用户在回调函数内可以读取输入 tensor 的有效地址，该指向前一层算子已经计算后的输出数据，输出 tensor 的有效地址指向即将送给下一层算子的输入。

rknn_custom_op_attr

开发者通过调用 rknn_custom_op_get_op_attr() 函数传入属性字段获得属性信息，属性信息用 rknn_custom_op_attr 表示，rknn_custom_op_attr 中的 void 类型 buffer，dtype 以及元素数量表示一块内存段，开发者根据 dtype 使用 C/C++ 将 buffer 强制转换指针类型可以得到相应数值类型的数组。

5.5.4.1.1 init 回调函数

常用于解析算子信息或初始化临时缓冲区或者输入/输出缓冲区 buffer。分配临时 buffer 的 init 回调函数示例代码如下：

  • CPU算子

/**
 * cpu kernel init callback for custom op
 */
int custom_op_init_callback(rkm_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs,
                          rknn_custom_op_tensor* outputs, uint32_t n_outputs)
{
    printf("custom_op_init_callback\n");
    // create tmp buffer
    float* tmp_buffer = (float*)malloc(inputs[0].attr.n_elems * sizeof(float));
    op_ctx->priv_data = tmp_buffer;
    return 0;
}

  • GPU算子

/**
 * opencnl kernel init callback for custom op
 **/
int relu_init_callback_gpu(rkm_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs,
                           rknn_custom_op_tensor* outputs, uint32_t n_outputs)
{
    printf("relu_init_callback_gpu\n");
    // 获取opecnl context
    cl_context cl_ctx = (cl_context)op_ctx->gpu_ctx.cl_context;

    // create tmp cl buffer
    cl_mem* memObject = (cl_mem*)malloc(sizeof(cl_mem) * 2);
    memObject[0] = clCreateBuffer(cl_ctx, CL_MEM_READ_WRITE, inputs[0].attr.size, NULL, NULL);
    memObject[1] = clCreateBuffer(cl_ctx, CL_MEM_READ_WRITE, outputs[0].attr.size, NULL, NULL);
    op_ctx->priv_data = memObject;
    return 0;
}

5.5.4.1.2 prepare 回调函数

该回调函数每帧推理都会调用，目前为预留实现。

5.5.4.1.3 compute 回调函数

它是自定义算子的计算函数，开发者必须完成输入/输出是 NCHW 或 UNDEFINED 格式 float32 数据类型输入输出的核函数。

1. compute 回调（CPU）

假设开发者想实现一个自定义层，完成 softmax 功能，CPU 算子 compute 函数示例如下：

/**
 * float32 kernel implemeatation sample for custom op
 **/
int compute_custom_softmax_float32(rkm_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs,
                                   rknn_custom_op_tensor* outputs, uint32_t n_outputs)
{
    unsigned char*  in_ptr  = (unsigned char*)inputs[0].mem.virt_addr + inputs[0].mem.offset;
    unsigned char*  out_ptr = (unsigned char*)outputs[0].mem.virt_addr + outputs[0].mem.offset;
    int             axis   = 0;
    const float*    in_data = (const float*)in_ptr;
    float*          out_data = (float*)out_ptr;
    std::string     name  = "";
    rknn_custom_op_attr op_name;
    rknn_custom_op_get_op_attr(op_ctx, "name", &op_name);
    if (op_name.n_elems > 0 && op_name.dtype == RKNN_TENSOR_UINT8) {
        name = (char*)op_name.data;
    }

    rknn_custom_op_attr op_attr;
    rknn_custom_op_get_op_attr(op_ctx, "axis", &op_attr);
    if (op_attr.n_elems == 1 && op_attr.dtype == RKNN_TENSOR_INT64) {
        axis = ((int64_t*)op_attr.data)[0];
    }

    printf("op name = %s, axis = %d\n", name.c_str(), axis);
    float* tmp_buffer = (float*)op_ctx->priv_data;
    // kernel implemeatation for custom op
    {
        int inside = 1;
        int outside = 1;
        int channel = 1;

        while (axis < 0) {
            axis += inputs[0].attr.n_dims;
        }
    }
    for (int i = 0; i < axis; i++) {
        outside *= inputs[0].attr.dims[i];
    }
    channel = inputs[0].attr.dims[axis];
    for (int i = axis; i < inputs[0].attr.n_dims; i++) {
        inside *= inputs[0].attr.dims[i];
    }

    for (int y = 0; y < outside; y++) {
        const float* src_y = in_data + y * inside;
        float*       dst_y = out_data + y * inside;
        float        max_data = -FLT_MAX;
        float        sum_data = 0.0f;

        for (int i = 0; i < inside; i++) {
            max_data = fmaxf(max_data, src_y[i]);
        }
        for (int i = 0; i < inside; i++) {
            tmp_buffer[i] = expf(src_y[i] - max_data);
            sum_data += tmp_buffer[i];
        }
        for (int i = 0; i < inside; i++) {
            dst_y[i] = tmp_buffer[i] / sum_data;
        }
    }
    return 0;
}

2. compute回调函数（GPU）

对于 GPU 算子，开发者可以在回调函数中完成以下步骤：

  • 开发者从 rkm_custom_op_context 里的 gpu_ctx 中获取 OpenCL 的 cl_context，cl_command_queue 以及 cl_kernel 对象，此过程需要开发者做数据类型转换。

  • 如有必要，用户自行创建的 op 输入或输出的 cl_mem 对象缓冲区。

  • 设置 cl_kernel 的函数参数。

  • OpenCL kernel 的函数参数的输入 buffer 数据目前只能支持 float，其他类型暂时还不支持。

  • 对于使用零拷贝的情况下，调用 clImportMemoryARM 可以自行协助用户把输入 tensor 的内存映射到 OpenCL 的 cl_mem 结构体中，输入 tensor 已包含输入数据，用户不需要自行再拷贝一次。该过程也可以在 init 回调函数中处理，然后将 cl_mem 结构体记录到 priv_data 成员，最后在 compute 回调中读取 priv_data 并使用它。

  • 以阻塞的形式运行 cl_kernel 。

  • CL kernel 内的输入数据都是以 NCHW 形式排布给出。

  • 如果在 GPU 运算完后，开发者需要 CPU 访问数据，需要通过调用 rkm_mem_sync 函数刷新输出 Tensor 的 cache 后再读取数据。

假设开发者想实现一个自定义层，完成relu功能，GPU算子compute函数示例如下：

/**
 * opencnl kernel init callback for custom op
 **/
int compute_custom_relu_float32(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t num_inputs,
                               rknn_custom_op_tensor* outputs, uint32_t num_outputs)
{
    std::string            name = "";
    rknn_custom_op_attr    op_name;
    rknn_custom_op_get_op_attr(op_ctx, "name", &op_name);
    if (op_name.n_elems > 0 && op_name.dtype == RKNN_TENSOR_UINT8) {
        name = (char*)op_name.data;
    }

    // get context
    cl_context cl_ctx = (cl_context)op_ctx->gpu_ctx.cl_context;

    // get command queue
    cl_command_queue queue = (cl_command_queue)op_ctx->gpu_ctx.cl_command_queue;

    // get kernel
    cl_kernel kernel = (cl_kernel)op_ctx->gpu_ctx.cl_kernel;

    // import input/output buffer
    const cl_import_properties_arm props[3] = {
    CL_IMPORT_TYPE_ARM,
    CL_IMPORT_TYPE_DMA_BUF_ARM,
    0,
};

cl_int status;
cl_mem inObject = clImportMemoryARM(cl_ctx, CL_MEM_READ_WRITE, props, &inputs[0].mem.fd,
                                    inputs[0].mem.offset + inputs[0].mem.size, &status);
if (status != CL_SUCCESS) {
    printf("Tensor: %s clImportMemoryARM failed\n", inputs[0].attr.name);
}

cl_mem outObject = clImportMemoryARM(cl_ctx, CL_MEM_READ_WRITE, props, &outputs[0].mem.fd,
                                     outputs[0].mem.offset + outputs[0].mem.size, &status);
if (status != CL_SUCCESS) {
    printf("Tensor: %s clImportMemoryARM failed\n", outputs[0].attr.name);
}

int     in_type_bytes = get_type_bytes(inputs[0].attr.type);
int     out_type_bytes = get_type_bytes(outputs[0].attr.type);
int     in_offset     = inputs[0].mem.offset / in_type_bytes;
int     out_offset    = outputs[0].mem.offset / out_type_bytes;
unsigned int elems    = inputs[0].attr.n_elems;

// set kernel args
int argIndex = 0;
clSetKernelArg(kernel, argIndex++, sizeof(cl_mem), &inObject);
clSetKernelArg(kernel, argIndex++, sizeof(cl_mem), &outObject);
clSetKernelArg(kernel, argIndex++, sizeof(int), &in_offset);
clSetKernelArg(kernel, argIndex++, sizeof(int), &out_offset);
clSetKernelArg(kernel, argIndex++, sizeof(unsigned int), &elems);

// set global worksize
const size_t global_work_size[3] = {elems, 1, 1};

    // enqueueNDRangeKernel
    clEnqueueNDRangeKernel(queue, kernel, 1, NULL, global_work_size, NULL, 0, NULL, NULL);

    // finish command queue
    clFinish(queue);

    // //cpu access data after sync to device
    // rkm_mem_sync(&outputs[0].mem, RKNN_MEMORY_SYNC_FROM_DEVICE);
    // // save output npy
    // char output_path[PATH_MAX];
    // sprintf(output_path, "%s/cpu_output%cd.npy", ".", 0);
    // unsigned char* out_data = (unsigned char*)outputs[0].mem.virt_addr + outputs[0].mem.offset;
    // save_npy(output_path, (float*)out_data, &inputs[0].attr);

    return 0;
}

5.5.4.1.4 destroy 回调函数

常用于销毁自定义算子的临时缓冲区或输入/输出 buffer。销毁临时 buffer 的示例代码如下：

  • CPU 算子

/**
 * cpu kernel destroy callback for custom op
 **/
int custom_op_destroy_callback(rkm_custom_op_context* op_ctx)
{
    printf("custom_op_destroy_callback\n");
    // clear tmp buffer
    free(op_ctx->priv_data);
    return 0;
}

  • GPU 算子

/**
 * opencnl kernel destroy callback for custom op
 **/
int relu_destroy_callback_gpu(rkm_custom_op_context* op_ctx)
{
    // clear tmp buffer
    printf("relu_destroy_callback_gpu\n");
    cl_mem* memObject = (cl_mem*)op_ctx->priv_data;
    clReleaseMemObject(memObject[0]);
    clReleaseMemObject(memObject[1]);
    free(memObject);
    return 0;
}

5.5.4.2 注册自定义算子

在设置完 rkm_custom_op 结构体后，需要调用 rkm_register_custom_ops() 将其注册到 rkm_context 中，该接口支持同时注册多个自定义算子。

在完成 CPU 的 compute 回调函数后，注册一个名为"cstSoftmax"和"ArgMax"的 CPU 自定义算子的代码示例如下：

  • CPU 算子

// CPU operators
rkm_custom_op user_op[2];
memset(user_op, 0, 2 * sizeof(rkm_custom_op));

strncpy(user_op[0].op_type, "cstSoftmax", RKNN_MAX_NAME_LEN - 1);
user_op[0].version = 1;
user_op[0].target = RKNN_TARGET_TYPE_CPU;
user_op[0].init = custom_op_init_callback;
user_op[0].compute = compute_custom_softmax_float32;
user_op[0].destroy = custom_op_destroy_callback;

strncpy(user_op[1].op_type, "ArgMax", RKNN_MAX_NAME_LEN - 1);
user_op[1].version = 1;
user_op[1].target = RKNN_TARGET_TYPE_CPU;
user_op[1].init = custom_op_init_callback;
user_op[1].compute = compute_custom_argmax_float32;
user_op[1].destroy = custom_op_destroy_callback;

ret = rkm_register_custom_ops(ctx, user_op, 2);
if (ret < 0) {
    printf("rkm_register_custom_ops fail! ret = %d\n", ret);
    return -1;
}

  • GPU算子

对于 GPU 算子而言，支持以常量字符串或者文件路径的两种方式注册 OpenCL kernel。当 rknn_custom_op 结构体中的 cl_source_size 等于 0 时，cl_kernel_source 表示 OpenCL kernel 的文件路径，当 cl_source_size 大于 0 时 cl_kernel_source 表示 OpenCL kernel 函数字符串。以字符串保存的 relu 功能的 OpenCL kernel 的示例代码如下：

char* cl_kernel_source = "#pragma OPENCL EXTENSION cl_arm_printf : enable"
    "#pragma OPENCL EXTENSION cl_khr_fp16 : enable "
    "__kernel void relu_float(__global const float* input, __global float* output, int in_offset, int out_offset, const unsigned int elems)"
    "{"
    "    int gid = get_global_id(0);"
    "    if (gid < elems) {"
    "        float in_value = input[in_offset + gid];"
    "        output[out_offset + gid] = in_value >= 0.f ? in_value : 0.f;"
    "    }"
    "}"
    "__kernel void relu_half(__global const half* input, __global half* output, int in_offset, int out_offset, const unsigned int elems)"
    "{"
    "    int gid = get_global_id(0);"
    "    if (gid < elems) {"
    "        half in_value = input[in_offset + gid];"
    "        output[out_offset + gid] = in_value >= 0.f ? in_value : 0.f;"
    "    }"
    "}";

在完成 OpenCL kernel 函数以及 GPU 的 compute 回调函数后，可以设置 rknn_custom_op 结构体数组并注册 GPU 算子，注册 GPU 算子示例代码如下：

// GPU operators
rknn_custom_op user_op[1];
memset(user_op, 0, sizeof(rknn_custom_op));
strncpy(user_op->op_type, "cstSoftmax", RKNN_MAX_NAME_LEN - 1);
user_op->version = 1;
user_op->target = RKNN_TARGET_TYPE_GPU;
user_op->init = relu_init_callback_gpu;
user_op->compute = compute_custom_relu_float32;
user_op->destroy = relu_destroy_callback_gpu;
#ifdef LOAD_FROM_PATH
user_op->cl_kernel_source = "/custom_op.cl";
user_op->cl_source_size = 0;
#else
user_op->cl_kernel_source = cl_kernel_source;
user_op->cl_source_size = strlen(cl_kernel_source);
#endif
strncpy(user_op->cl_kernel_name, "relu_float", RKNN_MAX_NAME_LEN - 1);

ret = rknn_register_custom_ops(ctx, user_op, 1);
if (ret < 0) {
    printf("rknn_register_custom_ops fail! ret = %d\n", ret);
    return -1;
}

注册调用该接口前要明确自定义算子的 op_type，准备好算子信息并配置 rkm_custom_op 数组。每个类型的自定义算子要调用一次注册接口，网络中同一类型的算子仅调用一次。

5.5.4.3 模型推理

在注册完所有算子后，可以使用通用 API 或零拷贝 API 流程完成推理。

5.5.4.4 连板精度分析

自定义算子的连板调试功能要求 rkm_server 版本 >=1.6.0。

连板调试时，RKNN Server 会采用 dlopen 的方式从特定目录打开用户编译好的自定义算子插件库来获取算子信息，对于插件库方式注册自定义算子，要求用户必须实现一个名为 get_rknn_custom_op 的函数。

若用户需要对包含自定义算子的模型做连板精度分析，具体步骤如下：

  1. 实现一个 get_rknn_custom_op() 函数和必须的回调函数，并编译成对应系统的库，编译的插件库名称必须以"librkscst_"为前缀，例如库名是 librkscst_relu.so。

  2. 插件放到 /vendor/lib64/（Android arm64-v8a）或 /usr/lib/rknpu/op_plugins（Linux）。

  3. 主机端或者上位机使用 RKNN-Toolkit2 的 Python 接口执行连板精度分析。

get_rkm_custom_op 函数的示例代码如下：

std::vector<std::string> get_all_plugin_paths(std::string plugin_dir)
{
    std::vector<std::string> plugin_paths;

    if (access(plugin_dir.c_str(), 0) != 0) {
        fprintf(stderr, "Can not access plugin directory: %s, please check it!\n", plugin_dir.c_str());
    }

    DIR* dir;
    struct dirent* ent;
    const char* prefix = RKNN_CSTOP_PLUGIN_PREFIX; // 所有库文件名应该以此前缀开头

    if ((dir = opendir(plugin_dir.c_str())) != NULL) {
        while ((ent = readdir(dir)) != NULL) {
            if (ent->d_type == DT_REG) {
                const char* filename = ent->d_name;
                size_t len = strlen(filename);

                if (len > 10 && strncmp(filename, prefix, strlen(prefix)) == 0) {
                    printf("Found plugin: %s file in %s\n", filename, plugin_dir.c_str());
                    plugin_paths.push_back(plugin_dir + "/" + filename);
                }
            }
        }
    


    closedir(dir);
} else {
    fprintf(stderr, "Unable to open directory");
}

return plugin_paths;
}

// the default path of the custom operator plugin libraries
std::string plugin_dir =
#if defined(__ANDROID__)
    #if defined(__aarch64__)
    "/vendor/lib64/";
    #else
    "/vendor/lib/";
    #endif // __aarch64__
#elif defined(__linux__)
    "/usr/lib/rknpu/op_plugins/";
#endif

std::vector<std::string> plugin_paths = get_all_plugin_paths(plugin_dir);
std::vector<void*> so_handles;
for (auto path : plugin_paths) {
    printf("load plugin %s\n", path.c_str());
    void* plugin_lib = dlopen(path.c_str(), RTLD_NOW);
    char* error = dlerror();
    if (error != NULL) {
        fprintf(stderr, "dlopen %s fail: %s.\nPlease try to set 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:%s'\n",
                path.c_str(), error, plugin_dir.c_str());
        dlclose(plugin_lib);
        return -1;
    }

printf("dlopen %s successfully!\n", path.c_str());
get_custom_op_func custom_op_func = (get_custom_op_func)dlsym(plugin_lib, "get_rkm_custom_op");
error = dlerror();
if (error != NULL) {
    fprintf(stderr, "dlsym fail: %s\n", error);
    dlclose(plugin_lib);
    return -1;
}

rkm_custom_op* user_op = custom_op_func();
ret = rkm_register_custom_ops(ctx, user_op, 1);
if (ret < 0) {
    printf("rkm_register_custom_ops fail! ret = %d\n", ret);
    return -1;
}
so_handles.push_back(plugin_lib);
}

插件库有如下注意事项：

  • 一个自定义算子插件库只能注册一个自定义算子，如果需要注册多个自定义算子，需要创建多个插件库。

  • 插件库的名称必须以librkcst_开头，以.so结尾。

  • 如果是C++代码实现的插件库，C接口中为了正确打开函数符号，要在函数定义前后加入如下宏：

#ifdef __cplusplus
extern "C" {
#endif

//code

#ifdef __cplusplus
} // extern "C"
#endif

  • 如果dlopen插件库失败，需要检查是否设置了 LD_LIBRARY_PATH 环境变量，并且检查插件库所在目录是否在环境变量指定的路径。

5.6 多Batch使用说明

5.6.1 多Batch原理

RK3588 NPU内部有3个核心，RK3576 NPU内部有2个核心，为了更高效得利用多核性能，提供了多batch推理功能。当开启多batch推理时，内部会调用 rknn_dup_context 将context进行拷贝（rknn_dup_context 只会对context的Internal进行拷贝，Weight会复用）。当 rknn_batch_size=2 时，会拷贝1份，当 rknn_batch_size >=3 时，会拷贝2份（同一时刻最多只有3个核心工作，为了避免内存浪费只拷贝2份）。每个context core_mask 会设置成0，让多核内部自动调度。当执行 rknn_run() 时，内部会起一个线程池，同一时刻调用3个线程同时对3个 context 进行推理。

图5-14 多batch内部原理图

5.6.2 多Batch使用方式

多batch使用方式如下：

  1. Python 端开启多 batch 设置：

ret = rkm.build(do_quantization=True, dataset='./dataset.txt', rknn_batch_size=3)

为了达到最优性能，RK3588建议 rknn_batch_size 为3的倍数，RK3576建议 rknn_batch_size 为2的倍数。

  1. 建议使用零拷贝接口

5.6.3 多Batch输入输出设置

当开启多 batch 功能时，查询出来的输入输出 size 是未开启时的 rknn_batch_size 倍。内部每个 context 会各自算自己的一个输入偏移量，按照这个输入偏移量取输入数据做推理，然后各自算自己的一个输出偏移量，按照这个输出偏移量写到各自的输出。以第二个 batch 为例，输入偏移量是查询出来的 input_size 除以 rknn_batch_size ，输出偏移量是查询出来的 output_size 除以 rknn_batch_size 。

图5-15 batch内部输入输出地址偏移图

5.7 RK3588 NPU SRAM使用说明

  • RK3588 SOC内部含有1MB的SRAM，其中有956KB可供给SOC上各个IP所使用。

  • SRAM可以帮助RKNP应用减轻DDR带宽压力，但对推理耗时可能有一定影响。

5.7.1 极端环境要求

5.7.1.1 内核环境要求

  • RKNPU驱动版本>=0.9.2

  • 内核config需要开启 CONFIG_ROCKCHIP_RKNPU_SRAM=y

    • Android系统config路径如下：

      <path-to-your-kernel>/arch/arm64/configs/rockchip_defconfig

    • Linux系统config路径如下：

      <path-to-your-kernel>/arch/arm64/configs/rockchip_linux_defconfig

  • 内核相应DTS需要从系统SRAM中分配给RKNPU使用

    • 从系统分配需求大小的SRAM给RKNPU，最大可分配956KB，且大小需要4K对齐。

    • 注意：默认系统中可能已为其他IP分配SRAM，比如编解码模块，各IP分配的SRAM区域不能重叠，否则会存在同时读写出现数据错乱现象。

    • 如下为956KB全部分配给RKNPU的例子：

syssram: sram@ff001000 {
    compatible = "mmio-sram";
    reg = <0x0 0xff001000 0x0 0x0ef000>;

    #address-cells = <1>;
    #size-cells = <1>;
    ranges = <0x0 0x0 0xff001000 0x0ef000>;
    /* 分配RKNPU SRAM / / start address and size should be 4k align */
    rknpu_sram: rknpu_sram@0 {
        reg = <0x0 0xef000>; // 956KB
    };
};

  • 把分配的SRAM挂到RKNPU节点，修改如下所示的dtsi文件：

<path-to-your-kernel>/arch/arm64/boot/dts/rockchip/rk3588s.dtsi 

rknpu: npu@fdab0000 {
    compatible = "rockchip,rk3588-rknpu";
    /* ... */ // 增加RKNPU sram的引用 
    rockchip,sram = <&rknpu_sram>;
    status = "disabled";
};

5.7.1.2 RKNN SDK版本要求

  • RKNPU Runtime库(librknnrt.so)版本>=1.6.0

5.7.2 使用方法

在 rknn_init() 接口的flags参数指定 RKNN_FLAG_ENABLE_SRAM 即可在该 context 中开启SRAM。

例如：

ret = rknn_init(&ctx, rknn_model, size, RKNN_FLAG_ENABLE_SRAM, NULL);

当设置RKNN_FLAG_ENABLE_SRAM时，将从系统可用的SRAM中分配尽可能多的内存做为模型的Internal Tensor内存。
注意：

  • 当SRAM被某一 rknn_context 占用后，其他的 rknn_context 不支持复用该段的SRAM。rknn_api.h 中的 RKNN_FLAG_SHARE_SRAM 功能暂未实现。

  • 当某个 rknn_context 未占用全部的SRAM时，剩余的SRAM可以给其他的 rknn_context 使用。

5.7.3 调试方法

5.7.3.1 SRAM是否启用查询

通过开机内核日志查看 SRAM 是否启用，包含为 RKNPU 指定 SRAM 的地址范围和大小信息，如下所示：

rk3588_s/# dmesg | grep rknpu -i
RKNPU f dab0000.npu: RKNPU: sram region: [0x00000000ff001000, 0x00000000ff0f0000), sram size: 0xef000

5.7.3.2 SRAM使用情况查询

  • 可通过节点查询SRAM的使用情况

  • 如下为未使用SRAM的位图表，每个点表示4K大小

rk3588_s/# cat /sys/kernel/debug/rknpu/mm
SRAM bitmap:"*" - used, "." - free (1bit = 4KB)
[000][.........................]
[001][.........................]
[002][.........................]
[003][.........................]
[004][.........................]
[005][.........................]
[006][.........................]
[007][.........]
SRAM total size: 978944, used: 0, free: 978944
# 单位为Byte

  • 如下为分配使用512KB后的SRAM位图表

rk3588_s/# cat /sys/kernel/debug/rknpu/mm
SRAM bitmap:"*" - used, "." - free (1bit = 4KB)
[000][*********************************]
[001][*********************************]
[002][*********************************]
[003][*********************************]
[004][.........................]
[005][.........................]
[006][.........................]
[007][.........]
SRAM total size: 978944, used: 524288, free: 454656
# 单位为Byte

5.7.3 通过RKNN API查询SRAM大小

  • 通过 rknn_query() 的 RKNN_QUERY_MEM_SIZE 接口查询SRAM大小信息

typedef struct _rknn_mem_size {
    uint32_t total_weight_size;
    uint32_t total_internal_size;
    uint64_t total_dma_allocated_size;
    uint32_t total_sram_size;
    uint32_t free_sram_size;
    uint32_t reserved[10];
} rknn_mem_size;

  • 其中，total_sram_size 表示：系统给RKNPU分配的SRAM总大小，单位是Byte。

  • free_sram_size 表示：剩余RKNPU能使用的SRAM大小，单位是Byte。

5.7.3.4 查看模型SRAM的占用情况

  • 板端环境中，RKNN应用运行前设置如下环境变量，可打印SRAM使用预测情况：

export RKNN_LOG_LEVEL=3

  • Internal分配SRAM的逐层占用情况，如下日志所示：

----------------------------------------------------------------------
Total allocated Internal SRAM Size: 524288, Addr: [0xff3e0000, 0xff460000)
----------------------------------------------------------------------
ID  User    Tensor  DataType  OrigShape        NativeShape       | [Start      End)       Size |  SramHit
----------------------------------------------------------------------
1   ConvRelu  input0  INT8    (1,3,224,224)  (1,224,224,3)  | 0xff3b0000 0xff3d4c00 0x00024c00 \ 
2   ConvRelu  output2 INT8    (1,32,112,112) (1,112,112,16) | 0xff404c00 0xff466c00 0x00062000 | 0x0005b400
3   ConvRelu  output4 INT8    (1,32,112,112) (1,4,112,112,16) | 0xff466c00 0xff52ac00 0x000c4000 | 0x00000000
4   ConvRelu  output6 INT8    (1,64,112,112) (1,4,112,112,16) | 0xff52ac00 0xff5ecc00 0x000c4000 | 0x00000000
5   ConvRelu  output8 INT8    (1,64,56,56)   (1,4,56,56,16) | 0xff5ecc00 0xff417000 0x00031000 | 0x00031000
6   ConvRelu  output10 INT8   (1,128,56,56)  (1,8,56,56,16) | 0xff411000 0xff473000 0x00062000 | 0x0004f000
7   ConvRelu  output12 INT8   (1,128,56,56)  (1,8,56,56,16) | 0xff473000 0xff4d5000 0x00062000 | 0x00000000
8   ConvRelu  output14 INT8   (1,128,56,56)  (1,8,56,56,16) | 0xff3ed000 0xff442000 0x00062000 | 0x00062000
9   ConvRelu  output16 INT8   (1,128,28,28)  (1,8,28,28,16) | 0xff442000 0xff45a800 0x00018800 | 0x00018800
10  ConvRelu  output18 INT8   (1,256,28,28)  (1,16,28,28,16) | 0xff3e0000 0xff411000 0x00031000 | 0x00031000
11  ConvRelu  output20 INT8   (1,256,28,28)  (1,16,28,28,16) | 0xff411000 0xff442000 0x00031000 | 0x00031000
12  ConvRelu  output22 INT8   (1,256,28,28) (1,16,28,28,16) | 0xff3e0000 0xff411000 | 0x00031000 | 0x00031000  
13  ConvRelu  output24 INT8   (1,256,14,14) (1,16,14,14,16) | 0xff411000 0xff41d400 | 0x0000c400 | 0x0000c400  
14  ConvRelu  output26 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3e0000 0xff3f8800 | 0x00018800 | 0x00018800 
15  ConvRelu  output28 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3f8800 0xff411000 | 0x00018800 | 0x00018800  
16  ConvRelu  output30 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3e0000 0xff3f8800 | 0x00018800 | 0x00018800  
17  ConvRelu  output32 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3f8800 0xff411000 | 0x00018800 | 0x00018800  
18  ConvRelu  output34 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3e0000 0xff3f8800 | 0x00018800 | 0x00018800  
19  ConvRelu  output36 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3f8800 0xff411000 | 0x00018800 | 0x00018800  
20  ConvRelu  output38 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3e0000 0xff3f8800 | 0x00018800 | 0x00018800  
21  ConvRelu  output40 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3f8800 0xff411000 | 0x00018800 | 0x00018800  
22  ConvRelu  output42 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3e0000 0xff3f8800 | 0x00018800 | 0x00018800  
23  ConvRelu  output44 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3f8800 0xff411000 | 0x00018800 | 0x00018800  
24  ConvRelu  output46 INT8   (1,512,14,14) (1,32,14,14,16) | 0xff3e0000 0xff3f8800 | 0x00018800 | 0x00018800  
25  ConvRelu  output48 INT8   (1,512,7,7) (1,33,7,7,16)      | 0xff3f8800 0xff3ff000 | 0x00006800 | 0x00006800  
26  ConvRelu  output50 INT8   (1,1024,7,7) (1,67,7,7,16)      | 0xff3e0000 0xff3ed000 | 0x0000d000 | 0x0000d000  
27  ConvRelu  output52 INT8   (1,1024,7,7) (1,67,7,7,16)      | 0xff3ed000 0xff3fa000 | 0x0000d000 | 0x0000d000  
28  AveragePool output54    INT8  (1,1024,7,7) (1,67,7,7,16)      | 0xff3e0000 0xff3ed000 | 0x0000d000 | 0x0000d000  
29  Conv        output55    INT8  (1,1024,1,1) (1,64,1,1,16)       | 0xff3ed000 0xff3ed400 | 0x00000400 | 0x00000400  
30  Softmax     output56    INT8  (1,1000,1,1) (1,64,1,1,16)       | 0xff3e0000 0xff3e0400 | 0x00000400 | 0x00000400  
31  OutputOperator output57 FLOAT (1,1000,1,1) (1,1000,1,1)        | 0xff3ae000 0xff3acf00 | 0x00000f00 | \           
-------------------------------------------+----------------------------------+-----------------
----------------------------------
Total Weight Memory Size: 4260864  
Total Internal Memory Size: 2157568  
Predict Internal Memory RW Amount: 11068320  
Predict Weight Memory RW Amount: 4260832  
Predict SRAM Hit RW Amount: 6688768  
----------------------------------

  • 其中上面文本图表中的SRAM Hit为当前层Tensor所占用的SRAM大小，表示与未开启SRAM对比，会节省对应大小的DDR读写数据量。

  • Predict SRAM Hit RW Amount为整个模型SRAM的读写预测情况，表示与未开启SRAM对比，每帧可节省的DDR读写数据量。

  • 注意：Linux环境日志重定向到终端，Android环境日志重定向到logcat。

5.8 模型剪枝

模型剪枝一般可分为有损剪枝和无损剪枝两种方式，RKNN-Toolkit2的模型剪枝为无损剪枝，也就是可以在减小模型大小和计算量的前提下，不会降低模型的浮点精度，甚至可以提高模型的量化精度。

但并不是所有模型都可以进行无损剪枝，无损剪枝是根据模型的权重的稀疏化程度，来去除一些对模型结果不造成影响的权重和Feature通道，以减小模型的大小和模型的计算量。在启用模型剪枝配置后（将 rknn.config() 的 model_pruning 参数设为 True ），模型转换时会自动根据权重的稀疏化程度对模型进行剪枝。

如果模型剪枝成功，则会打印剪枝结果，如下：

I model_pruning results:
I Weight: -1.12145 MB (-6.9%)
I GFLOPs: -0.15563 (-13.4%)

其中，Weight剪掉了1.12145 MB（占比6.9%），模型运算量减少0.15563 GFLOPs（占比13.4%）。

如果因模型稀疏化程度不够而剪枝失败，则不会做任何处理（也不会打印上述信息），因此该配置并不会影响正常的模型转换。

5.9 模型加密

模型加密指的是生成完RKNN模型后，再重新对其做进一步处理。使用 rknn.export_rknn() 生成的模型，可以通过Netron等第三方工具查看图结构。模型加密后，Netron等第三方工具将无法查看相应的网络结构，也无法获取权重，起到对模型的保护作用。当前加密后的RKNN模型使用方法和未加密的模型一样，不需要在开发板推理时做任何修改。

使用方法如下：

# Create RKNN object
rknn = RKNN()
# Export encrypted RKNN model
crypt_level = 1
ret = rknn.export_encrypted_rknn_model("input.rknn", "encrypt.rknn", crypt_level)
if ret != 0:
    print("Encrypt RKNN model failed!")

crypt_level 用来指定加密等级，有1，2和3三个等级。默认值为1。等级越高，安全性越高，解密越耗时；反之，安全性越低，解密越快。

  • 支持平台：RK3562/RK3566/RK3568/RK3576/RK3588

5.10 Cacheable内存一致性

Cacheable内存一致性问题是指当CPU和NPU设备都会访问同一块带cache标志的内存时，CPU会将数据缓存到cache中，如果NPU访问到的DDR数据与CPU cache不一致，会导致读取数据错误。因此要调用刷新cache的接口保证CPU和NPU访问到的DDR内存数据是一致的。本章节介绍了同步数据的方向以及如何使用 rknn_mem_sync 接口刷新cache。

5.10.1 Cacheable内存同步的方向

当CPU写数据到cacheable的内存，之后NPU访问该内存时，要保证CPU cache的数据同步到DDR中，此时同步的方向是指从CPU到NPU设备；当NPU写完数据，CPU开始访问该内存时，要保证DDR的数据与CPU cache中的一致，此时同步的方向是指从NPU设备到CPU。RKNN C API提供了 rknn_mem_sync_mode 枚举类型表示cacheable内存同步的方向，数据结构如下：

/*
The mode to sync cacheable rknn memory.
*/
typedef enum _rknn_mem_sync_mode {

RKN_MEMORY_SYNC_TO_DEVICE = 0x1, /* the mode used for consistency of device access after CPU accesses data. */

RKN_MEMORY_SYNC_FROM_DEVICE = 0x2, /* the mode used for consistency of CPU access after device accesses data. */

RKN_MEMORY_SYNC_BIDIRECTIONAL =
RKN_MEMORY_SYNC_TO_DEVICE | RKN_MEMORY_SYNC_FROM_DEVICE, /* the mode used for consistency of data access
between device and CPU in both directions. */

} rknn_mem_sync_mode;

  • RKN_MEMORY_SYNC_TO_DEVICE：表示数据同步方向是CPU到NPU设备

  • RKN_MEMORY_SYNC_FROM_DEVICE：表示数据同步方向是NPU设备到CPU

  • RKN_MEMORY_SYNC_BIDIRECTIONAL：表示数据在NPU和CPU之间双向同步，在用户不确定同步数据的方向时可以使用该枚举

在明确数据同步方向时，建议使用单方向刷新CPU cache模式，能避免多余的刷新CPU cache动作导致性能损耗。

5.10.2 同步Cacheable内存

明确数据同步方向之后，就可以使用 rknn_mem_sync 接口对cacheable内存做同步。接口形式如下：

int rknn_mem_sync(rknn_context context, rknn_tensor_mem* mem, rknn_mem_sync_mode mode);

其中，context 是上下文（在非RV1103系列/RV1106系列平台上默认设置为NULL），mem 是 rknn_create_mem 接口返回的 rknn_tensor_mem* 指针类型，mode 是指同步数据的方向。

5.11 模型稀疏化推理

该功能用于Torch模型在训练阶段对模型权重自动稀疏化并使用RKNN进行稀疏化推理，提高模型推理速度。目前该功能仅支持RK3576。

5.11.1 稀疏化原理

模型权重稀疏化在训练阶段根据用户自定义的方式，进行权重置零操作，具体方式有4:2输入方向稀疏化，4:2输出方向稀疏化，16:4输入输出稀疏化，16:4输出输入稀疏化。其中4:2输入方向稀疏化实现原理如图5-16所示，将模型权重沿输入方向，在连续的4个数值中选择两个置零。需要注意的是，当模型权重指定方向非4对齐时，会进行补齐操作。

图5-16 4:2输入方向稀疏化原理图 有关四种权重稀疏化说明如表5-16所示，其中输入输出方向为权重的方向，例如2维卷积shape为 ( C_{out} \times C_{in} \times K_h \times K_w )，输入方向即 ( C_{in} )，输出方向即 ( C_{out} ) 。在稀疏率为75%的方式中前置训练模型可以是未稀疏化模型也可以是稀疏模型，建议使用单方向稀疏模型作为前置预训练模型，可减少稀疏化后的精度损失。 表5-10 模型权重稀疏化说明

稀疏化方式

稀疏率

说明

4:2输入方向稀疏化

50%

沿输入方向对权重稀疏化

4:2输出方向稀疏化

50%

沿输出方向对权重稀疏化

16:4输入输出稀疏化

75%

先沿输入方向4:2稀疏，再沿输出方向4:2稀疏

16:4输出输入稀疏化

75%

先沿输出方向4:2稀疏，再沿输入方向4:2稀疏

5.11.2 训练稀疏化模型

  1. 首先确保cuda可使用并安装对应python版本的autosparsity包

pip install autosparsity-1.0-cp38-cp38m-linux_x86_64.whl

autosparsity安装包路径为：https://github.com/airockchip/rknn-toolkit2/tree/master/autosparsity/packages

  1. 以torchvision中的resnet50的4:2 输入方向稀疏化为例，进行权重稀疏化训练

import torch
import torchvision.models as models
from autosparsity.sparsity import sparsity_model

if __name__ == "__main__":
    model = models.resnet50(pretrained=True).cuda()
    optimizer = None
    mode = 0
    sparsity_model(model, optimizer, mode)
    model.eval()
    x = torch.randn(1, 3, 224, 224).cuda()
    torch.onnx.export(
        model, x,'resnet50.onnx', input_names=['inputs'], output_names=['outputs']
    )

自定义模型的稀疏化在模型训练之前添加 sparsity_model 函数即可，参考示例如下：

# insert model autosparsity code before training
import torch
import torchvision.models as models
from autosparsity.sparsity import sparsity_model

...

model = models.resnet34(pretrained=True).cuda()
mode = 0
sparsity_model(model, optimizer, mode)

# normal training
x, y = DataLoader(args)
for epoch in range(epochs):
    y_pred = model(x)
    loss = loss_func(y_pred, y)
    loss.backward()

    optimizer.step()
...

注：Attention模型稀疏化训练依赖于一个已训练好的前置模型。所以在做稀疏化训练之前，请先训练一个效果“不错”的模型，在该模型基础上进行稀疏化训练。

有关 sparsity_model 函数的参数说明如下： 表5-11 sparsity_model函数各参数说明

参数

详细说明

model

原训练模型

optimizer

原优化器，默认为None

mode

稀疏化方式，可选值为0,1,2,3，默认为0：0:4:2输入方向稀疏化(50%稀疏率) 1:4:2输出方向稀疏化(50%稀疏率) 2:16:4输入输出稀疏化(75%稀疏率) 3:16:4输出输入稀疏化(75%稀疏率)

verbose

log等级，可选值为0,1,2,3，默认为2：0: Errors 1: Errors and Warnings 2: Errors, warnings and info

whitelist

稀疏化支持的module列表，支持1d conv, 2d conv，3d conv, linear, MultiheadAttention，默认[torch.nn.Linear, torch.nn.Conv2d]

allowed_layer_names

允许稀疏化的层名，用户配置时则只稀疏指定层，默认None

disallowed_layer_names

不允许稀疏化的层名，用户配置时则会跳过该层，默认[]

fast

设为True代表使用快速方法计算mask，默认为False(默认的mask计算方法针对部分模型会失效，若稀疏化报错可尝试将该参数设为True)

5.11.3 RKNN稀疏化推理使用方法

通过RKNN-Toolkit2中 config() 接口的"sparse_infer"参数设置模型稀疏化的开启和关闭，对应的参数为True/False，默认值为 False 。开启稀疏化推理的参考代码如下：

rknn.config(target_platform='rk3576', sparse_infer=True)

完整的稀疏化Python推理代码可参考：https://github.com/airockchip/rknn-toolkit2/tree/master/autosparsity/examples

使用C API进行部署时，首先使用RKNN-Toolkit2在 config() 接口中设置"sparse_infer"参数为True生成带稀疏化推理的RKNN模型，之后正常调用通用API接口流程或零拷贝接口流程即可。

使用Python代码推理时可在构建RKNN对象时设置verbose=True，开启日志打印各层稀疏化情况；使用C API推理时通过设置环境变量 RKNN_LOG_LEVEL=4，开启日志打印各层稀疏化情况。日志信息如下：

图5-17 4:2输入方向稀疏化Python打印日志

![](images/RKNN_SDK_高级用户指南/4:2输入方向稀疏化C API打印日志.png) 图5-17 4:2输入方向稀疏化Python打印日志

其中 SparseRation 代表稀疏率，4:2输入方向稀疏化对应50%(IC)，4:2输出方向稀疏化对应50%(OC)，16:4输入输出稀疏化和16:4输出输入稀疏化对应75%。0%则代表未做稀疏化。

5.11.4 RKNN稀疏化推理限制

稀疏化推理是基于NPU的硬件架构实现，受硬件规格限制，RK3576的稀疏化推理目前只支持2d Conv并且group参数为1，其余限制如下：

表5-12 RK3576稀疏化推理限制

通道数量 数据类型
4:2输入方向稀疏化 输入32对齐，输出32对齐 Int8 Float16不支持
4:2输出方向稀疏化 输入32对齐，输出32对齐 Int8 Float16不支持
16:4输入输出稀疏化 输入32对齐，输出32对齐 Int8 Float16不支持
16:4输出输入稀疏化 输入32对齐，输出32对齐 Int8 Float16不支持

5.12 生成部署C代码

RKNN-Toolkit2 2.0.0版本新增 Codegen 接口，用于生成模型部署代码，简化开发者的上手难度。Codegen 基于 CAPI 零拷贝接口进行二次封装，接口风格与 RKNN_Model_Zoo 中的 demo 一致。生成的部署代码可用于直接测试性能、验证精度，开发者也可以基于生成的代码进行二次开发。

使用示例：

ret = rknn.codegen(output_path='./rknn_app_demo',
                   inputs=['../../caffe/mobilenet_v2/dog_224x224.jpg'], overwrite=True)

  • 调用 codegen 接口前，必须先调用 rknn.export 接口保存 RKNN 模型。

  • output_path 为输出文件夹目录，用户可配置目录名称。

  • inputs 填写模型输入的路径列表，允许不填。有效文件格式为 .jpg / .png / .npy ，以 .npy 文件为输入时，.npy 数据的维度信息应与模型输入的维度信息保持一致。

  • overwrite 设为 True 时，会覆盖 output_path 指定目录下的文件。默认值为 False 。

  • 生成部署代码后，请参考生成目录下的 README.md 文档说明进行编译、测试。

  • 若 inputs 填入有效值，部署代码示例在推理后，会评估 CAPI 接口与 RKNN-Toolkit2 模拟器之间的推理结果差异，评估方式为对比每一个输出的余弦相似度。

  • 无 NPU 硬件平台限制，要求板端系统为 Linux 或 Android。

  • 支持量化、非量化模型。

5.13 ONNX模型编辑

部分模型在转RKNN模型后，可能存在冗余的op，常见于模型的输入输出节点处存在冗余的reshape、transpose op，影响了RKNN模型的推理性能。RKNN-Toolkit2提供了 onnx_edit 接口，用于修改ONNX模型的输入输出的维度定义，使调整后的onnx模型能转出性能更好的RKNN模型，减少冗余的reshape、transpose op。

5.13.1 onnx_edit接口说明

使用示例：

from rknn.utils import onnx_edit
ret = onnx_edit(model='./concat_block.onnx',
                export_path='./concat_block_edited.onnx',
                inputs_transform={'k_in': 'a,b,c,d->1,a,d,b,c'},
                outputs_transform={'k_cache': 'a,b,c,d->1,a,b,c,d'},
                dataset='./dataset.txt'
)

  • model : 填入待修改模型，为必填参数

  • export_path : 填入新模型的生成路径，为必填参数

  • inputs_transform : 填入输入节点的变换公式字典，key 为节点名称，value 为变换公式。为可选参数，默认为空字典

  • outputs_transform : 填入输出节点的变换公式字典。为可选参数，默认为空字典

  • dataset : 填入输入数据的路径集文件，文件格式要求与 rknn.build 接口对 dataset 的要求一致。填入后，onnx_edit 接口除了对模型的输入输出定义进行调整，也会将 dataset 中对应的数据进行调整，在 export_path 同级目录下生成新的 dataset 数据，可用于新模型的验证、量化。为可选参数，默认为空。

5.13.2 onnx_edit变换公式说明

onnx_edit 接口中的 inputs_transform、outputs_transform 需要填入变换公式。变换公式的定义与 einsum 算子的定义类似，例如 'a,b,c,d->1,a,d,b,c' 公式，指将原始的维度为 'a,b,c,d'，变换后为 '1,a,d,b,c'。变换公式的填写规则如下：

  • 必须有两部分字符组成，并用 '->' 隔开，左边为原始字符 shape，右边是变换后字符 shape

  • 字符 shape 里面的符号只允许是 [a-z] 或 ','、'1'，除了 '1' 以外，不支持其他数字 shape 字符，原因是相同数字字符无法判断 transpose 的前后关系

  • 字符 [a-z] 认为是独立的，没有顺序关系，即 'a,b->b,a' 和 'c,a->a,c' 表示的是一样的变换

  • 等式左边的原始字符 shape，被 ',' 分隔成 n 份，n 的个数必须和模型的维度匹配，例如模型输入定义是 [32,4,1,64]，输入字符可以是 'a,b,c,d' 或 'a,b,1,d'

  • 原始字符 shape 的每一个符号，除了 '1'，都必须存在于变换后字符 shape，例如 'a,1,c,d->a,c,d,1' 是有效的，'a,1,c,d->a,c,1' 是无效的，缺了 'd'

  • 变换后字符 shape，允许插入任意个数的 '1'，达成扩维效果，例如 'a,b,c->a,1,c,b,1,1'

  • 原始字符 shape，允许用多个字母以及赋值公式来表示对 shape 进行拆分，例如原始输入定义是 [32,4,1,64]，'ab,c,d,qk[a=2,k=8]->aq,cd,1,kb'，表示将 32 拆分成 2x16，将 64 拆分成 8x8，再进行 transpose, reshape 操作。其中 '[]' 的部分称为赋值公式，多个公式用 ',' 符号分隔。此外，允许拆分中的某个字符没有赋值，此时会自动推断对应的 shape，例如赋值公式只给了 a=2，已知在模型中 ab=32，则自动推断出 b=16；若推断出的 shape 异常会直接报错，比如 ab=32，若赋值 a=5，则 b=6.4，又维度必须是整数，此时会抛出异常错误。

5.13.3 变换公式示例

  • 将3维输入修改为4维输入：'a,b,c->a,b,1,c'

  • 将5维输入修改为4维输入：'a,b,c,d,e->ab,c,d,e'

  • 进行transpose(0,3,1,2)操作：'a,b,c,d->a,d,b,c'

  • Transpose并合并部分维度：'a,b,c,d->d,acb,1'

  • 拆分维度、transpose、合并维度：'a,bc,de,fl[b=2,d=4]->ab,fe,dc,1'