[RKNN] 1. 入门介绍
[RKNN] 2. 模型转换和推理–API介绍&以yolox为例
[RKNN] 3. 零拷贝接口推理
[RKNN] 4. 基于零拷贝接口封装
第一篇文章已经环境都配置好了,并且进行了测试。接下来就需要进行自己的工作了,由于之前做过yolox在TensorRT、NCNN的部署,为了更好的对比不同框架的表现,所以本文也以yolox为例。本文主要介绍RKNN的模型转换,主要是从onnx转成rknn,从文档里复制一些常用API,并实现yolox的模型转换和推理。
首先应该知道虽然rknn也是类似和tensorrt一样与实际硬件相关的,但是模型的转换过程官方只给了在上位机(非开发板)中的接口,所以模型转换部分均要在上位机(一般是安装ubuntu的x86系统上,也可以是wsl,我用的就是wsl),并且只有python的转换接口。
模型初始化
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()
在上位机也可以通过连接板子或者利用库自带的模拟器,进行推理测试。
运行环境初始化
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,它可以在摄像机演示中使用
)
RKNN-Toolkit2 1.4.0 及之前的版本支持1.6.0~1.9.0 版本ONNX 导出的模型。onnx版本和opset的对应可以参考网站,1.9.0最高支持opset14.我在实验中发现如果用opset14,后续转rknn的时候报错,提示建议opset12。
在训练好yolox的模型后,可以参考我的代码进行导出。注意我修改了yolox的模型结构,focus
模块换成了Conv
,silu
换成了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()
然后在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的推理。
rknpu的C接口比较全面,整体可以分为两组,通用API接口和零拷贝流程的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;
}
}
这个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);
首先应该说明,我个人现阶段不建议使用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()
C++的接口就很全了,直接推理的话可以参考我写的下面这个程序。tools.hpp
和postprocess.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=(Dq−ZP)∗Sq
static float deqnt_affine_to_f32(int8_t qnt, int32_t zp, float scale) {
return ((float)qnt - (float)zp) * scale;
}