[RKNN] 2. 模型转换和推理--API介绍&以yolox为例

系列文章目录

[RKNN] 1. 入门介绍
[RKNN] 2. 模型转换和推理–API介绍&以yolox为例
[RKNN] 3. 零拷贝接口推理
[RKNN] 4. 基于零拷贝接口封装


文章目录

  • 系列文章目录
  • 前言
  • 一、模型转换
    • 1.1 API介绍
      • 1.1.1 初始化&加载&构建&导出模型
      • 1.1.2 模型测试
    • 2. yolox模型转换
      • 1.2.1导出onnx
      • 1.2.2 onnx2rknn
  • 二、模型推理
    • 2.1 C API介绍
      • 2.1.1 推理API
      • 2.1.2 矩阵乘法API
    • 2.2 yolox模型推理
      • 2.2.1 Python
      • 2.2.2 CPP
  • 附录 量化


前言

第一篇文章已经环境都配置好了,并且进行了测试。接下来就需要进行自己的工作了,由于之前做过yolox在TensorRT、NCNN的部署,为了更好的对比不同框架的表现,所以本文也以yolox为例。本文主要介绍RKNN的模型转换,主要是从onnx转成rknn,从文档里复制一些常用API,并实现yolox的模型转换和推理。


一、模型转换

首先应该知道虽然rknn也是类似和tensorrt一样与实际硬件相关的,但是模型的转换过程官方只给了在上位机(非开发板)中的接口,所以模型转换部分均要在上位机(一般是安装ubuntu的x86系统上,也可以是wsl,我用的就是wsl),并且只有python的转换接口。

1.1 API介绍

1.1.1 初始化&加载&构建&导出模型

模型初始化

from rknn.api import RKNN
# 初始化
rknn = RKNN(verbose=True, verbose_file=None)
# verbose表示是否打印日志信息 verbose_file表示日志文件

# 模型配置 返回0 / -1,0表示成功
# 主要修改的是均值方差和目标平台
rknn.config(
    mean_values: Any | None = None,		# 输入均值
    std_values: Any | None = None,		# 输入方差
    quantized_dtype: str = 'asymmetric_quantized-8',	# 量化数据类型
    quantized_algorithm: str = 'normal',	# 量化算法,normal mmse kl_divergence
    quantized_method: str = 'channel',		# 量化方法,layer channel
    target_platform: Any | None = None,		# 目标平台,我的是rk3588s
    quant_img_RGB2BGR: bool = False,		# 量化图像是否rgb2bgr
    float_dtype: str = 'float16',			# 非量化的数据类型,只支持fp16
    optimization_level: int = 3,			# 设置优化级别,3是最高
    custom_string: Any | None = None,		# 将字符串信息添加到rknn模型中,运行时会打印
    remove_weight: bool = False,			# 去除conv2d的权重,
    compress_weight: bool = False,			# 压缩权重,可以减小文件大小
    inputs_yuv_fmt: Any | None = None,		# 增加yuv预处理
    single_core_mode: bool = False,			# 减小rknn模型大小,只适用rk3588
    dynamic_input: Any | None = None,		# 动态输入:[[[1,3,224,224],[1,1,224,224], ...], [[1,3,160,160],[1,1,160,160], ...], ...]
    model_pruning: bool = False,			# 剪枝
    op_target: Any | None = None,			# 给每个操作设置设备:{'111':'cpu', '222':'cpu', ...}
    **kwargs: Any
) 

模型加载。支持onnx、pytorch、tensorflow等模型的权重。我这里主要列出pytorch和onnx的api。

# onnx  返回0 / -1,0表示成功
# 主要指定model就可以了
rknn.load_onnx(
    model: Any,		# onnx 文件路径
    inputs: Any | None = None,				# 指定模型输入。None表示从模型获取
    input_size_list: Any | None = None,		# 设置每个张量的大小列表
    input_initial_val: Any | None = None,	# 设置每个输入的初始值
    outputs: Any | None = None				# 设置模型输出。None表示从模型获取
)

# pytorch 返回0 / -1,0表示成功
# 主要指定model就可以了
rknn.load_pytorch(
    model: Any,				# pytorch权重文件
    input_size_list: Any	# 设置每个输入的大小 [[1,224,224],[3,224,224]]
)

模型构建

# 模型构建 返回0 / -1,0表示成功
 rknn.build(
    do_quantization: bool = True,		# 是否量化
    dataset: Any | None = None,			# 数据集
    rknn_batch_size: Any | None = None	# 推理batch_size
)

模型导出

# 模型导出 返回0 / -1,0表示成功
rknn.export_rknn(
    export_path: Any,				# 导出路径
    cpp_gen_cfg: Any | None = None,	# 是否生成c++部署实例
    **kwargs: Any
)

模型销毁

rknn.release()

1.1.2 模型测试

在上位机也可以通过连接板子或者利用库自带的模拟器,进行推理测试。
运行环境初始化

rknn.init_runtime(
    target: Any | None = None,			# 目标平台。None的时候会运行仿真环境
    target_sub_class: Any | None = None,# ?
    device_id: Any | None = None,		# adb的id
    perf_debug: bool = False,			# 是否打开性能调试
    eval_mem: bool = False,				# 是否打开内存调试
    async_mode: bool = False,			# 是否启动异步模式
    core_mask: int = NPU_CORE_AUTO		# 指定NPU的核
    '''
    NPU_CORE_AUTO  = 0                                   # default, run on NPU core randomly.
    NPU_CORE_0     = 1                                   # run on NPU core 0.
    NPU_CORE_1     = 2                                   # run on NPU core 1.
    NPU_CORE_2     = 4                                   # run on NPU core 2.
    NPU_CORE_0_1   = 3                                   # run on NPU core 1 and core 2.
    NPU_CORE_0_1_2 = 7                                   # run on NPU core 1 and core 2 and core 3.
    '''
) 

推理

rknn.inference(
    inputs: Any,		# 输入的np数组
    data_format: Any | None = None,			# 输入类型‘nhwc’ 'nchw' 
    inputs_pass_through: Any | None = None,	# 将输入透传给 NPU 驱动。非透传模式下,在将输入传给 NPU 驱动之前,工具会对输入进行减均值、除方差等操作;而透传模式下,不会做这些操作。
    get_frame_id: bool = False 	# 当使用异步模式时是否需要获取输出/输入帧id,它可以在摄像机演示中使用
)

2. yolox模型转换

1.2.1导出onnx

RKNN-Toolkit2 1.4.0 及之前的版本支持1.6.0~1.9.0 版本ONNX 导出的模型。onnx版本和opset的对应可以参考网站,1.9.0最高支持opset14.我在实验中发现如果用opset14,后续转rknn的时候报错,提示建议opset12。
在训练好yolox的模型后,可以参考我的代码进行导出。注意我修改了yolox的模型结构,focus模块换成了Convsilu换成了relu

from loguru import logger

import torch
from torch import nn
from torch import Tensor
from yolox.exp import get_exp

class Config:
    name        = "yolo-s"  # 类型
    exp_file    = "../yolox_s.py" # 配置文件 需要修改
    ckpt_file   = './YOLOX_outputs/yolox_s/best_ckpt.pth' # 权重文件 需要修改

    output_name = "yolox_relu_no_decode.onnx"  # 输出onnx文件
    input       = "images"      # onnx input name
    output      = "output"      # onnx ouput name
    opset       = 12            # op 版本
    batch_size  = 1             # .
    dynamic     = False          # 动态尺寸
    onnxsim     = True          # 简化onnx

    opts = None # Modify config options using the command-line
    decode_in_inference = False # decoder

@logger.catch
def export_zzx():
    exp = get_exp(Config.exp_file, Config.name)
    Config.experiment_name = exp.exp_name
    model = exp.get_model()
  
    # load the model state dict
    ckpt = torch.load(Config.ckpt_file, map_location="cpu")

    model.eval()
    if "model" in ckpt:
        ckpt = ckpt["model"]
    model.load_state_dict(ckpt)
    model.head.decode_in_inference = Config.decode_in_inference

    logger.info("loading checkpoint done.")
    dummy_input = torch.randn(Config.batch_size, 3, exp.test_size[0], exp.test_size[1])

    torch.onnx._export(
        model,
        dummy_input,
        Config.output_name,
        input_names=[Config.input],
        output_names=[Config.output],
        dynamic_axes={Config.input: {0: 'batch'},
                      Config.output: {0: 'batch'}} if Config.dynamic else None,
        opset_version=Config.opset,
    )
    logger.info(f"generated onnx model named {Config.output_name}")

    if Config.onnxsim:
        import onnx
        from onnxsim import simplify

        # use onnx-simplifier to reduce reduent model.
        onnx_model = onnx.load(Config.output_name)
        model_simp, check = simplify(onnx_model)
        assert check, "Simplified ONNX model could not be validated"
        onnx.save(model_simp, Config.output_name)
        logger.info(f"generated simplified onnx model named {Config.output_name}")

if __name__ == "__main__":
    export_zzx()

1.2.2 onnx2rknn

然后在rknn的环境中,就可以把onnx转成rknn了,我的芯片是rk3588s。

import numpy as np
import cv2

from rknn.api import RKNN

ONNX_MODEL = 'onnx/yolox_relu_no_decode.onnx'
RKNN_MODEL = 'yolox_relu_nodecode.rknn'
IMG_PATH = 'inference/1.jpg'
DATASET = './dataset.txt'	# 里面保存量化需要的图片路径,应该多一点并且贴近后续实际输入
DEVICE_NAME = 'rk3588s'

QUANTIZE_ON = True
IMG_SIZE = 640

def convert():
    rknn = RKNN(verbose=True, verbose_file=None)
    
    # pre-process config
    rknn.config(mean_values=[[0, 0, 0]], std_values=[[1, 1, 1]], target_platform=DEVICE_NAME, quant_img_RGB2BGR = True)

    # Load ONNX model
    ret = rknn.load_onnx(model=ONNX_MODEL)
    if ret != 0:
        print('Load model failed!')
        exit(ret)

    # Build model
    ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
    if ret != 0:
        print('Build model failed!')
        exit(ret)

    # Export RKNN model

    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    rknn.release()

if __name__ == '__main__':
    convert()

如果想要顺便仿真运行一遍可以在rknn.release()之前一行加上下面的代码

ret = rknn.init_runtime()
if ret != 0:
    print('Init runtime environment failed!')
    exit(ret)

# Set inputs
img = cv2.imread(IMG_PATH)
img, ratio = preproc(img, (IMG_SIZE, IMG_SIZE))	# 预处理 根据自己的训练调整后

# Inference
outputs = rknn.inference(inputs=[img[None, :, :, :]], data_format='nchw')

# 后处理 主要包括decode和nms。
dets = postprocess(outputs[0], (IMG_SIZE, IMG_SIZE))	

二、模型推理

上面的整个过程都是在上位机中完成,在实际中一般需要在开发板上运行和部署,这里我也进行利用瑞芯微提供的开发板推理接口进行yolox的推理。

2.1 C API介绍

rknpu的C接口比较全面,整体可以分为两组,通用API接口和零拷贝流程的API接口。

  • 通用接口:每次更新数据,需要把数据拷贝到NPU运行的输入内存。
  • 零拷贝接口:使用预先分配的内存,减少内存拷贝的花销。
    注意两者不能混合调用。其中通用接口就类似GPU编程,申请提前显存。零拷贝接口有三种方式,包括输入/输出内存由运行时分配、输入/输出由外部分配、输入/输出/权重/中间结果内存由外部分配。

[RKNN] 2. 模型转换和推理--API介绍&以yolox为例_第1张图片[RKNN] 2. 模型转换和推理--API介绍&以yolox为例_第2张图片

[RKNN] 2. 模型转换和推理--API介绍&以yolox为例_第3张图片[RKNN] 2. 模型转换和推理--API介绍&以yolox为例_第4张图片
值得注意的是,零拷贝API需要满足的条件:

  1. 输入通道数是1 或3 或4。
  2. RK3588输入宽度需要16像素对齐
  3. int8是非对称量化

2.1.1 推理API

rknn_init: 初始化rknn

rknn_context ctx;
int ret = rknn_init(&ctx, model_data, model_data_size, 0, NULL);
// 1 rknn_context 指针
// 2 RKNN 模型的二进制数据或者RKNN 模型路径
// 3 当model 是二进制数据,表示模型大小,当model 是路径,则设置为0
// 4 flag.	RKNN_FLAG_COLLECT_PERF_MASK:用于运行时查询网络各层时间;
//			RKNN_FLAG_MEM_ALLOC_OUTSIDE:用于表示模型输入、输出、权重、中间tensor 内存全部由用户分配;
//			RKNN_FLAG_SHARE_WEIGHT_MEM:用于表示共享另个模型的weight 权重;
//			RKNN_FLAG_COLLECT_MODEL_INFO_ONLY:用于初始化一个空上下文,它只可以调用rknn_query 接口查询模型weight 内存总大小和中间tensor 总大小,而无法进行推理。
// 5 rknn_init_extend 特定初始化时的扩展信息。没有使用,传入NULL 即可。如果需要共享模型weight 内存,则需要传入另个模型rknn_context 指针。

rknn_set_core_mask: 设置NPU核心(只支持3588/3588s)

rknn_context ctx;
rknn_core_mask core_mask = RKNN_NPU_CORE_0;
int ret = rknn_set_core_mask(ctx, core_mask);
//	RKNN_NPU_CORE_AUTO:表示自动调度模型,自动运行在当前空闲的NPU 核上;
//	RKNN_NPU_CORE_0:表示运行在NPU0 核上;
//	RKNN_NPU_CORE_1:表示运行在NPU1 核上;
//	RKNN_NPU_CORE_2:表示运行在NPU2 核上;
//	RKNN_NPU_CORE_0_1:表示同时工作在NPU0、NPU1 核上;
//	RKNN_NPU_CORE_0_1_2:表示同时工作在NPU0、NPU1、NPU2 核上。

rknn_dup_context: 生成指向同一个模型的新context,可用于多线程执行相同模型时的权重复用.(好东西)

rknn_context ctx_in;	// old
rknn_context ctx_out;	// new
int ret = rknn_dup_context(&ctx_in, &ctx_out);

rknn_destroy: 释放传入的rknn_context 及其相关资源。(放到析构里面)

rknn_context ctx;
int ret = rknn_destroy (ctx);

rknn_query: 查询获取到模型输入输出信息、逐层运行时间、模型推理的总时间、SDK 版本、内存占用信息、用户自定义字符串等信息。(啥啥都能查)

rknn_context ctx;
int ret = rknn_quert(ctx, 查询命令, 返回的数据, 数据结构的大小);
// 查询命令					数据结构					功能
// RKNN_QUERY_IN_OUT_NUM 	rknn_input_output_num 	查询输入输出tensor个数
// RKNN_QUERY_INPUT_ATTR 	rknn_tensor_attr 		查询输入tensor属性
// RKNN_QUERY_OUTPUT_ATTR 	rknn_tensor_attr 		查询输出tensor属性
// RKNN_QUERY_PERF_DETAIL 	rknn_perf_detail 		查询网络各层运行时间,需要rknn_init设置RKNN_FLAG_COLLECT_PERF_MASK
// RKNN_QUERY_PERF_RUN 		rknn_perf_run 			查询推理模型(不包含设置输入/输出)耗时,微秒
// RKNN_QUERY_SDK_VERSION 	rknn_sdk_version 		查询SDK版本
// RKNN_QUERY_MEM_SIZE 		rknn_mem_size 			查询分配给权重和网络中间tensor的内存大小
// RKNN_QUERY_CUSTOM_STRING rknn_custom_string 		查询RKNN模型里面的用户自定义字符串信息

// RKNN_QUERY_NATIVE_INPUT_ATTR  rknn_tensor_attr 	使用零拷贝API时,查询原生输入tensor属性,它是NPU直接读取的模型输入属性
// RKNN_QUERY_NATIVE_OUTPUT_ATTR rknn_tensor_attr 	使用零拷贝API时,查询原生输出tensor属性,它是NPU直接输出的模型输出属性
// RKNN_QUERY_NATIVE_NC1HWC2_INPUT_ATTR  rknn_tensor_attr 使用零拷贝API时,查询原生输入tensor属性,它是NPU直接读取的模型输入属性与RKNN_QUERY_NATIVE_INPUT_ATTR查询结果一致
// RKNN_QUERY_NATIVE_NC1HWC2_OUTPUT_ATTR rknn_tensor_attr 使用零拷贝API时,查询原生输出tensor属性,它是NPU直接输出的模型输出属与RKNN_QUERY_NATIVE_OUTPUT_ATTR查询结果一致
// RKNN_QUERY_NATIVE_NHWC_INPUT_ATTR 	 rknn_tensor_attr 使用零拷贝API时,查询原生输入tensor属性与RKNN_QUERY_NATIVE_INPUT_ATTR查询结果一致
// RKNN_QUERY_NATIVE_NHWC_OUTPUT_ATTR    rknn_tensor_attr 使用零拷贝API时,查询原生输出NHWC tensor属性

// RKNN_QUERY_INPUT_DYNAMIC_RANGE 	rknn_input_range 使用支持动态形状RKNN模型时,查询模型支持输入形状数量、列表、形状对应的数据布局和名称等信息
// RKNN_QUERY_CURRENT_INPUT_ATTR 	rknn_tensor_attr 使用支持动态形状RKNN模型时,查询模型当前推理所使用的输入属性
// RKNN_QUERY_CURRENT_OUTPUT_ATTR 	rknn_tensor_attr 使用支持动态形状RKNN模型时,查询模型当前推理所使用的输出属性

rknn_inputs_set: 设置模型的输入数据。该函数能够支持多个输入,其中每个输入是rknn_input 结构体对象,在传入之前用户需要设置该对象。

rknn_input inputs[1];
memset(inputs, 0, sizeof(inputs));
inputs[0].index = 0;
inputs[0].type = RKNN_TENSOR_UINT8;
inputs[0].size = img_width*img_height*img_channels;
inputs[0].fmt = RKNN_TENSOR_NHWC;
inputs[0].buf = in_data;
inputs[0].pass_through = 0;		// 0数据根据fmt和type进行变换然后进行计算。 1 则不进行变换
ret = rknn_inputs_set(ctx, 1, inputs);	// 1表示输入数据个数

rknn_run: 执行一次推理,调用前需要通过rknn_inputs_set或者零拷贝接口设置输入数据。

ret = rknn_run(ctx, NULL);

rknn_outputs_get: 获取推理输出结果。可以获取多个输出数据,每个数据都是rknn_output结构体对象,需要在调用函数之前创建和设置。
输出数据的存放可以采用两种方式:一种是用户自行申请和释放,此时rknn_output 对象的is_prealloc 需要设置为1,并且将buf 指针指向用户申请的buffer;另一种是由rknn 来进行分配,此时rknn_output 对象的is_prealloc 设置为0 即可,函数执行之后buf将指向输出数据.

rknn_output outputs[io_num.n_output];
memset(outputs, 0, sizeof(outputs));
for (int i = 0; i < io_num.n_output; i++) {
outputs[i].index = i;
outputs[i].is_prealloc = 0;
outputs[i].want_float = 1;
} ret =
rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);
// ctx	输入个数	数组数据保存的数组	扩展

rknn_outputs_release: 释放rknn_outputs_get得到的输出相关资源。

ret = rknn_outputs_release(ctx, io_num.n_output, outputs);
// ctx	输入个数	数组数据保存的数组

rknn_create_mem: 用户需要自己分配内存让NPU使用,可以利用该函数创建一个rknn_tensot_mem结构体,并得到其指针。函数参数包括内存大小,返回值是结构体指针。

rknn_tensor_mem* input_mems[1];
input_mems[0] = rknn_create_mem(ctx, size);

rknn_create_mem_from_phys: 用户需要自己分配内存让NPU使用时,可以利用该函数创建一个rknn_tensot_mem结构体,并得到其指针。函数参数包括物理地址、虚拟地址、内存大小,返回值是结构体指针。

rknn_tensor_mem* input_mems[1];
input_mems[0] = rknn_create_mem_from_phys(ctx, input_phys, input_virt, size);

rknn_create_mem_from_fd: 用户需要自己分配内存让NPU使用,可以利用该函数创建一个rknn_tensot_mem结构体,并得到其指针。函数参数包括文件描述符、偏移、虚拟地址、内存大小,返回值是结构体指针。

rknn_tensor_mem* input_mems[1];
input_mems[0] = rknn_create_mem_from_fd(ctx, input_fd, input_virt, size, 0);

rknn_destroy_mem: 销毁rknn_tensor_mem结构体。但是我们申请的保存数据的内存还是需要额外进行释放。

rknn_tensor_mem* input_mems [1];
int ret = rknn_destroy_mem(ctx, input_mems[0]);

rknn_set_weight_mem: 实现用户为网络权重分配内存,初始化rknn_tensor_mem结构体后,rknn_run之前,利用该函数可以让NPU使用内存。

rknn_tensor_mem* weight_mems[1];
int ret = rknn_set_weight_mem(ctx, weight_mems[0]);

rknn_set_internal_mem: 实现用户为网络中间tensor分配内存,初始化rknn_tensor_mem结构体后,rknn_run之前,利用该函数可以让NPU使用内存。

rknn_tensor_mem* internal_tensor_mems [1];
int ret = rknn_set_internal_mem(ctx, internal_tensor_mems[0]);

rknn_set_io_mem: 实现用户为网络输入/输出tensor分配内存,初始化rknn_tensor_mem结构体后,rknn_run之前,利用该函数可以让NPU使用内存。

rknn_tensor_attr output_attrs[1];	// 输出信息
rknn_tensor_mem* output_mems[1];	// 输出tensor
ret = rknn_query(ctx, RKNN_QUERY_NATIVE_OUTPUT_ATTR, &(output_attrs[0]),
sizeof(rknn_tensor_attr));		// 查询信息
output_mems[0] = rknn_create_mem(ctx, output_attrs[0].size_with_stride); // 分配内存
rknn_set_io_mem(ctx, output_mems[0], &output_attrs[0]);	// 设置

rknn_set_input_shape: 对于动态形状输入RKNN 模型,在推理前必须指定当前使用的输入形状。该接口传入的rknn_tensor_attr*参数包含了输入形状和对应的数据布局信息,将rknn_tensor_attr 结构体对象的dims 成员设置输入数据的形状,将fmt 成员要设置成对应的数据布局信息。在使用该接口前,可先通过rknn_query 查询RKNN 模型支持的输入数量和动态形状列表,要求输入数据的形状在模型支持的输入形状列表中。每次切换新的输入形状,需要调用该接口设置新的形状。

for (int i = 0; i < io_num.n_input; i++) {
	for (int j = 0; j < input_attrs[i].n_dims; ++j) {
		//读取第0 个动态输入形状
		input_attrs[i].dims[j] = dyn_range[i].dyn_range[0][j];
	}
	ret = rknn_set_input_shape(ctx, &input_attrs[i]);
	if (ret < 0) {
		fprintf(stderr, "rknn_set_input_shape error! ret=%d\n", ret);
	return -1;
	}
}

2.1.2 矩阵乘法API

这个API可以实现矩阵的乘法在NPU上的加速执行。这个API接收两个二维矩阵输入返回其乘积矩阵。 C = A × B C=A \times B C=A×B,其中 A A A M M M K K K列, B B B K K K N N N列, C C C M M M N N N列。
注意: 受限于NPU硬件,对于矩阵乘法也由一些限制,对于int8输入矩阵,要求K小于等于4096;对于float16输入矩阵,要求K小于等于2048。其次,K和N均只能大于32,且是32的倍数。输入只支持8-bit有符号整型和16-bit 浮点,输出只支持32-bit有符号整型和32-bit浮点。

rknn_matmul_create: 创建矩阵乘法运算句柄,上下文的初始化,同时读取矩阵乘法规格信息,获取输入和输出tensor 属性。rknn_matmul_ctx 句柄与rknn_context 是同一个数据结构.

rknn_matmul_info info;	// 设置信息
memset(&info, 0, sizeof(rknn_matmul_info));
info.M = 4;
info.K = 64;
info.N = 32;
info.type = RKNN_TENSOR_INT8;
info.native_layout = 0;
info.perf_layout = 0;
rknn_matmul_io_attr io_attr;
memset(&io_attr, 0, sizeof(rknn_matmul_io_attr));
int ret = rknn_matmul_create(&ctx, &info, &io_attr);

rknn_matmul_set_io_mem: 设置矩阵乘法运算的输入/输出内存。在调用该函数前,先使用rknn_create_mem接口创建的rknn_tensor_mem结构体指针,接着将其与rknn_matmul_create 函数返回的矩阵A、B或C 的rknn_matmul_tensor_attr 结构体指针传入该函数,把输入和输出内存设置到矩阵乘法上下文中。在调用该函数前,要根据rknn_matmul_info 中配置的内存排布准备好矩阵A 和矩阵B的数据。

rknn_tensor_mem* A = rknn_create_mem(ctx, io_attr.A.size);
memset(A->virt_addr, 1, A->size);
rknn_matmul_io_attr io_attr;
memset(&io_attr, 0, sizeof(rknn_matmul_io_attr));
int ret = rknn_matmul_create(&ctx, &info, &io_attr);
// Set A
ret = rknn_matmul_set_io_mem(ctx, A, &io_attr.A);

rknn_matmul_set_core_mask: 设置矩阵乘法运算时可用的NPU核心(仅支持RK3588)。在调用前,需要通过rknn_matmul_create 函数初始化矩阵乘法上下文。可通过该函数设置的掩码值,指定需要使用的核心,以提高矩阵乘法运算的性能和效率。

rknn_matmul_set_core_mask(ctx, RKNN_NPU_CORE_AUTO);

rknn_matmul_run: 运行矩阵乘法运算,结果保存在输出矩阵C中。在调用函数前,输入矩阵A和B需要先准备好数据,并通过rknn_matmul_set_io_mem 函数设置到输入缓冲区。输出矩阵C需要先通过rknn_matmul_set_io_mem 函数设置到输出缓冲区,而输出矩阵的tensor 属性则通过
rknn_matmul_create 函数获取。

int ret = rknn_matmul_run(ctx);

rknn_matmul_destroy: 销毁矩阵乘法运算上下文,释放相关资源。在使用完rknn_matmul_create 函数创建的矩阵乘法运算句柄后,需要调用该函数进行销毁。

int ret = rknn_matmul_destroy(ctx);

2.2 yolox模型推理

2.2.1 Python

首先应该说明,我个人现阶段不建议使用rknn_toolkit_lite2进行部署(仅指当前1.5版本)。板载python推理接口只能推理,没有测试性能之类的接口,不过既然是做部署的,兄弟们也迟早得写C++,建议跳过python。
下面这个程序是我推理的一个demo,demo_postprocess, multiclass_nms这三个方法可以见yolox的官方仓库。

import cv2
import time
import numpy as np
from rknnlite.api import RKNNLite

from tools import demo_postprocess, multiclass_nms

# RKNN_MODEL = 'yolox_relu_decode.rknn'
RKNN_MODEL = 'yolox_relu_nodecode.rknn'
IMG = '1.jpg'
IMG_SIZE = (640, 640)
NMS_THR = 0.65
CON_THR = 0.45

# 图像预处理
def preproc(img: np.ndarray, input_size: tuple, swap: tuple=(2, 0, 1))->Tuple[np.ndarray, float]:
    padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
    
    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
    resized_img = cv2.resize(img, 
                             (int(img.shape[1] * r), int(img.shape[0] * r)),
                             interpolation=cv2.INTER_LINEAR,
                             ).astype(np.uint8)
    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img

    # padded_img = padded_img.transpose(swap)
    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
    return padded_img, r

def yolox_infer():
    # 初始化
    rknn_lite = RKNNLite()  
    # 加载模型
    ret = rknn_lite.load_rknn(RKNN_MODEL)
    if ret != 0:
        print('Load RKNN model failed')
        exit(ret)
    
    # 初始化运行环境 设置单npu核
    ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
    if ret != 0:
        print('Init runtime environment failed')
        exit(ret)

    # 加载图片
    ori_img = cv2.imread(IMG)

    t1 = time.perf_counter()
    ## 图像预处理
    img, ratio = preproc(ori_img, IMG_SIZE)
    # cv2.imwrite("a.jpg", img)

    # 推理
    outputs = rknn_lite.inference(inputs=[img[None, :, :, :]])

    # 后处理
    predictions = np.squeeze(outputs[0])
    predictions = demo_postprocess(predictions, IMG_SIZE)

    # nms
    boxes = predictions[:, :4]
    scores = predictions[:, 4:5] * predictions[:, 5:]

    boxes_xyxy = np.ones_like(boxes)
    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
    boxes_xyxy /= ratio
    dets = multiclass_nms(boxes_xyxy, scores, NMS_THR, CON_THR)

    t2 = time.perf_counter()
    print(f"infer time: {(t2 - t1) * 1000}ms")

    print(f"infer ans: {dets}")

    # 销毁
    rknn_lite.release()

if __name__ == '__main__':
    yolox_infer()

2.2.2 CPP

C++的接口就很全了,直接推理的话可以参考我写的下面这个程序。tools.hpppostprocess.hpp是两个功能函数和它们的实现,都粘贴过来太长了,可以参考我的Github。实际上下面这个程序已经完成了利用rknn api推理的全流程。但是可以改进的地方还很多比如类封装、用板载RGA进行预处理、零拷贝等等,后面会继续写的。

#include 
#include 
#include 

#include "rga.h"
#include "im2d.h"
#include "rknn_api.h"
#include "opencv2/opencv.hpp"

#include "tools.hpp"
#include "postprocess.hpp"

// 打印信息
static void dump_tensor_attr(rknn_tensor_attr* attr){

	std::string shape_str = attr->n_dims < 1 ? "" : std::to_string(attr->dims[0]);
	for (int i = 1; i < attr->n_dims; ++i) {
		shape_str += ", " + std::to_string(attr->dims[i]);
	}

	printf("  index=%d, name=%s, n_dims=%d, dims=[%s], n_elems=%d, size=%d, w_stride = %d, size_with_stride=%d, fmt=%s, "
			"type=%s, qnt_type=%s, "
			"zp=%d, scale=%f\n",
			attr->index, attr->name, attr->n_dims, shape_str.c_str(), attr->n_elems, attr->size, attr->w_stride,
			attr->size_with_stride, get_format_string(attr->fmt), get_type_string(attr->type),
			get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}

// resize img
cv::Mat static_resize(cv::Mat& img, int input_w, int input_h) {
    float r = std::min(input_w / (img.cols*1.0), input_h / (img.rows*1.0));
    // r = std::min(r, 1.0f);
    int unpad_w = r * img.cols;
    int unpad_h = r * img.rows;
    cv::Mat re(unpad_h, unpad_w, CV_8UC3);
    cv::resize(img, re, re.size());
    cv::Mat out(input_h, input_w, CV_8UC3, cv::Scalar(114, 114, 114));
    re.copyTo(out(cv::Rect(0, 0, re.cols, re.rows)));
    return out;
}


int main(){

	std::string model_name = "/home/orangepi/zzx/yolox_rknn/3infer_cpp/yolox_relu_nodecode.rknn";
	std::string image_name = "../img/1.jpg";
	const float nms_threshold      = 0.65;
	const float box_conf_threshold = 0.45;

	rknn_context   ctx;
	std::vector<rknn_tensor_attr> input_attrs;
    std::vector<rknn_tensor_attr> output_attrs;

	// 反量化参数
	std::vector<float>    out_scales;
	std::vector<int32_t>  out_zps;

	// 加载文件
	int model_data_size = 0;
	unsigned char* model_data = load_model(model_name.c_str(), &model_data_size);

	// 初始化
	CHECK_RKNN(rknn_init(&ctx, model_data, model_data_size, 0, NULL));
	
	// 获取&打印版本信息
	rknn_sdk_version version;
	CHECK_RKNN(rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version)));
	printf("sdk version: %s driver version: %s\n", version.api_version, version.drv_version);

	// 获取&打印输入输出数量
	rknn_input_output_num io_num;
	CHECK_RKNN(rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num)));
	printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);

	// 获取&打印input信息
	input_attrs.resize(io_num.n_input);
	// memset(input_attrs, 0, sizeof(input_attrs));
	for (int i = 0; i < io_num.n_input; i++) {
		input_attrs[i].index = i;
		CHECK_RKNN(rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr)));
		printf("input information: ");
		dump_tensor_attr(&(input_attrs[i]));
	}
	// 获取&打印output信息
	output_attrs.resize(io_num.n_output);
	// memset(output_attrs, 0, sizeof(output_attrs));
	for (int i = 0; i < io_num.n_output; i++) {
		output_attrs[i].index = i;
		CHECK_RKNN(rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr)));
		printf("output information:");
		dump_tensor_attr(&(output_attrs[i]));
	}

	int channel = 3;
	int width   = 0;
	int height  = 0;
	if (input_attrs[0].fmt == RKNN_TENSOR_NCHW) {
		printf("model is NCHW input fmt\n");
		channel = input_attrs[0].dims[1];
		height  = input_attrs[0].dims[2];
		width   = input_attrs[0].dims[3];
	} else {
		printf("model is NHWC input fmt\n");
		height  = input_attrs[0].dims[1];
		width   = input_attrs[0].dims[2];
		channel = input_attrs[0].dims[3];
	}

	// 初始化输入
	rknn_input inputs[1];
	memset(inputs, 0, sizeof(inputs));
	inputs[0].index        = 0;
	inputs[0].type         = RKNN_TENSOR_UINT8;
	inputs[0].size         = width * height * channel;
	inputs[0].fmt          = RKNN_TENSOR_NHWC;				
	inputs[0].pass_through = 0;

	// 申请输出空间
	rknn_output outputs[io_num.n_output];
	memset(outputs, 0, sizeof(outputs));
	for (int i = 0; i < io_num.n_output; i++) {
		outputs[i].want_float = false;			// 输出是u8类型。 true则在内部转成fp后再输出	
	}

	// 初始化后处理类
	for (int i = 0; i < io_num.n_output; ++i) {
		out_scales.push_back(output_attrs[i].scale);
		out_zps.push_back(output_attrs[i].zp);
	}
	std::shared_ptr<YoloxPostProcess> post_process = std::make_shared<YoloxPostProcess>(height, box_conf_threshold, nms_threshold, output_attrs);

	// 读取图片
	cv::Mat img = cv::imread(image_name, 1);
	// cv::cvtColor(orig_img, img, cv::COLOR_BGR2RGB);

	// 预处理
	float scale = std::min(width / (img.cols*1.0), height / (img.rows*1.0));
	auto img_out = static_resize(img, width, height);
	inputs[0].buf = (void*)img_out.data;

	// 设置输入
	CHECK_RKNN(rknn_inputs_set(ctx, io_num.n_input, inputs));

	// 推理
	CHECK_RKNN(rknn_run(ctx, NULL));

	// 获取输出
	CHECK_RKNN(rknn_outputs_get(ctx, io_num.n_output, outputs, NULL));

	// 后处理
	auto res = post_process->process((int8_t *)outputs->buf, out_zps, out_scales);
	
	// 打印结果
	printf("res size: %ld\n", res.size());
	for (auto a : res) {
		std::cout << scale << std::endl;
		a.x1 /= scale;
		a.y1 /= scale;
		a.x2 /= scale;
		a.y2 /= scale;

		std::cout<<a.x1<<" "<<a.y1<<" "<<a.x2<<" "<<a.y2 <<" "<<a.score<<" "<<a.category<< std::endl; 
		cv::rectangle(img, cv::Point(a.x1, a.y1), cv::Point(a.x2, a.y2), cv::Scalar(255, 0, 0, 255), 3);
		cv::putText(img, std::to_string(a.category), cv::Point(a.x1, a.y1 + 12), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));
	}
	cv::imwrite("./out.jpg", img);
	CHECK_RKNN(rknn_outputs_release(ctx, io_num.n_output, outputs));

	// 测速
	int test_count = 200;
	// warmup
	for (int i = 0; i < 50; ++i) {
		auto img_out = static_resize(img, width, height);
		inputs[0].buf = (void*)img_out.data;
		CHECK_RKNN(rknn_inputs_set(ctx, io_num.n_input, inputs));
		CHECK_RKNN(rknn_run(ctx, NULL));
		CHECK_RKNN(rknn_outputs_get(ctx, io_num.n_output, outputs, NULL));
		auto res = post_process->process((int8_t *)outputs->buf, out_zps, out_scales);
		CHECK_RKNN(rknn_outputs_release(ctx, io_num.n_output, outputs));
	}
	
	auto start = std::chrono::system_clock::now();
	for (int i = 0; i < test_count; ++i) {
		auto img_out = static_resize(img, width, height);
		inputs[0].buf = (void*)img_out.data;
		CHECK_RKNN(rknn_inputs_set(ctx, io_num.n_input, inputs));
		CHECK_RKNN(rknn_run(ctx, NULL));
		CHECK_RKNN(rknn_outputs_get(ctx, io_num.n_output, outputs, NULL));
		auto res = post_process->process((int8_t *)outputs->buf, out_zps, out_scales);
		CHECK_RKNN(rknn_outputs_release(ctx, io_num.n_output, outputs));
	}

	auto end = std::chrono::system_clock::now();
	float infer_time = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() ;
	printf("运行 %d 次, 平均耗时 %f ms\n", test_count, infer_time / (float)test_count);

	// release
	CHECK_RKNN(rknn_destroy(ctx));

	if (model_data) {
		free(model_data);
	}
	return 0;
}

附录 量化

量化和反量化的方法和参数可以通过rknn_query接口查询。

std::vector<rknn_tensor_attr> input_attrs.resize(io_num.n_input);
std::vector<rknn_tensor_attr> output_attrs.resize(io_num.n_output);
rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr))

float32->int8(非对称量化),假设输入张量的非对称量化参数是 S q S_q Sq, Z P ZP ZP,数据 D D D的量化过程可以表示为:
D q = r o u n d ( c l a m p ( D / S q + Z P ) , − 128 , 127 ) D_q=round(clamp(D/S_q+ZP),-128, 127) Dq=round(clamp(D/Sq+ZP),128,127)
float32->int16(非对称量化),假设输入张量的非对称量化参数是 S q S_q Sq, Z P ZP ZP,数据 D D D的量化过程可以表示为:
D q = r o u n d ( c l a m p ( D / S q + Z P ) , − 32768 , 32767 ) D_q=round(clamp(D/S_q+ZP),-32768, 32767) Dq=round(clamp(D/Sq+ZP),32768,32767)

通过指定rknn_output outputs[0].want_float = false;可以设置我们需要输出是否是float类型,如果设置为false则需要手动进行反量化。
int8->float32: D = ( D q − Z P ) ∗ S q D=(D_q-ZP)*S_q D=(DqZP)Sq

static float deqnt_affine_to_f32(int8_t qnt, int32_t zp, float scale) { 
    return ((float)qnt - (float)zp) * scale; 
}

你可能感兴趣的:(RKNN,计算机视觉,边缘计算,人工智能)