当前位置：首页 > news >正文

rknn优化教程（三）

news 来源：原创 2025/6/23 5:43:54

文章目录

- 1. 前述
- 2. 部分代码
- 3. 说明

1. 前述

OK，这一篇博客将完整给出最后的优化教程，包括代码设计。

首先有这样的目录结构：

./rknn_engine
├── include
│   ├── def
│   │   └── rknn_define.h
│   └── rknn_engine.h
├── src
│   ├── common
│   │   ├── rknn_data.h
│   │   └── rknn_functions.hpp
│   ├── inference
│   │   ├── inference.cpp
│   │   └── inference.h
│   ├── postprocess
│   │   ├── postprocess.cpp
│   │   └── postprocess.h
│   ├── preprocess
│   │   ├── preprocess.cpp
│   │   └── preprocess.h
│   ├── rknn_engine.cpp
│   └── task
│       ├── base_task.h
│       ├── pool_task.cpp
│       ├── pool_task.h
│       ├── single_task.cpp
│       ├── single_task.h
│       └── task_define.h
├── xmake.lua
└── xmake_repo10 directories, 18 files

其实这里只给出了detection的部分设计，其他的segment和pose就是对应扩充一下就可以了。我觉得还是很简单的……

2. 部分代码

就给出一些代码示意吧：

Inference::init

bool Inference::init(const rknn_binary &model_data){int ret = rknn_init(&m_rknnCtx, const_cast<char *>(&model_data[0]), static_cast<uint32_t>(model_data.size()), 0, nullptr);if (ret < 0){spdlog::error("RKNN初始化失败, ret={}", ret);return false;}m_isInited = true;// 配置运行在什么核心上ret = rknn_set_core_mask(m_rknnCtx, rk3588_npu[m_ctxIndex++ % NPU_NUMS]);if (ret != RKNN_SUCC){spdlog::error("rknn_set_core_mask failed, ret:{}", ret);return false;}// 获取版本信息rknn_sdk_version version;ret = rknn_query(m_rknnCtx, RKNN_QUERY_SDK_VERSION, &version, sizeof(rknn_sdk_version));if (ret < 0){spdlog::error("RKNN查询版本失败, ret={}", ret);return false;}spdlog::info("RKNN API version: {}", version.api_version);spdlog::info("RKNN Driver version: {}", version.drv_version);// 获取输入输出的个数rknn_input_output_num io_num;ret = rknn_query(m_rknnCtx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));if (ret != RKNN_SUCC){spdlog::error("RKNN查询输入输出失败, ret={}", ret);return false;}spdlog::info("模型的输入数量: {}, 输出数量: {}", io_num.n_input, io_num.n_output);// 模型的输入属性m_modelParams.m_nInput = io_num.n_input;for (uint32_t index = 0; index < io_num.n_input; ++index){rknn_tensor_attr attr{0};attr.index = index;ret = rknn_query(m_rknnCtx, RKNN_QUERY_INPUT_ATTR, &attr, sizeof(attr));if (ret != RKNN_SUCC){spdlog::error("RKNN查询输入属性失败, ret={}", ret);return false;}logTensorAttr(attr);m_modelParams.m_inputAttrs.push_back(attr);}// 模型的输出属性m_modelParams.m_nOutput = io_num.n_output;for (uint32_t index = 0; index < io_num.n_output; ++index){rknn_tensor_attr attr{0};attr.index = index;ret = rknn_query(m_rknnCtx, RKNN_QUERY_OUTPUT_ATTR, &attr, sizeof(attr));if (ret != RKNN_SUCC){spdlog::error("RKNN查询输出属性失败, ret={}", ret);return false;}logTensorAttr(attr);m_modelParams.m_outputAttrs.push_back(attr);}// 判断是否是量化的auto &out1_attr = m_modelParams.m_outputAttrs[0];if (out1_attr.qnt_type == RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC && out1_attr.type == RKNN_TENSOR_INT8){m_modelParams.m_isFloat = false;}else{m_modelParams.m_isFloat = true;}// 获得宽高和通道auto &in1_attr = m_modelParams.m_inputAttrs[0];if (in1_attr.fmt == RKNN_TENSOR_NCHW){spdlog::info("model is NCHW input fmt.");m_modelParams.m_modelChannel = in1_attr.dims[1];m_modelParams.m_modelHeight = in1_attr.dims[2];m_modelParams.m_modelWidth = in1_attr.dims[3];}else{spdlog::info("model is NHWC input fmt.");m_modelParams.m_modelChannel = in1_attr.dims[3];m_modelParams.m_modelHeight = in1_attr.dims[1];m_modelParams.m_modelWidth = in1_attr.dims[2];}spdlog::info("model input height:{} width:{} channel:{}",m_modelParams.m_modelHeight, m_modelParams.m_modelWidth, m_modelParams.m_modelChannel);spdlog::info("RKNN初始化成功!");return true;}

SingleTask::work

bool SingleTask::work(const cv::Mat &img, std::vector<DetectionResult> &dets){ 前处理 //m_preprocess.run(img, m_pImgBuff); 推理 //// inputrknn_input inputs[m_modelParams.m_nInput];memset(inputs, 0, sizeof(inputs));inputs[0].index = 0;inputs[0].type = RKNN_TENSOR_UINT8;inputs[0].fmt = RKNN_TENSOR_NHWC;inputs[0].size = m_imgSize;inputs[0].buf = m_pImgBuff;// outputrknn_output outputs[m_modelParams.m_nOutput];memset(outputs, 0, sizeof(outputs));for (uint32_t index = 0; index < m_modelParams.m_nOutput; ++index){outputs[index].index = index;outputs[index].want_float = m_modelParams.m_isFloat;}m_inference.run(inputs, outputs); 后处理 //m_postprocess->run(outputs);int width = img.cols;int height = img.rows;auto scale = m_preprocess.getScale();auto &labels = m_postprocess->getLabels();auto &results = m_postprocess->detectionResult();int label_count = static_cast<int>(labels.size());for (auto &result : results){DetectionResult rd;int id = result.m_classId < 0 ? -1 : (result.m_classId < label_count ? result.m_classId : -1);rd.m_className = id < 0 ? "unknown" : labels[id];rd.m_confidence = result.m_confidence;rd.m_box.x = static_cast<int>(result.m_x * scale);rd.m_box.x = clamp_i(rd.m_box.x, 0, width);rd.m_box.y = static_cast<int>(result.m_y * scale);rd.m_box.y = clamp_i(rd.m_box.y, 0, height);int w = static_cast<int>(result.m_w * scale);int h = static_cast<int>(result.m_h * scale);rd.m_box.w = clamp_i(w, 0, width - rd.m_box.x);rd.m_box.h = clamp_i(h, 0, height - rd.m_box.y);dets.push_back(rd);}// release outputrknn_outputs_release(m_inference.rknnContext(), m_modelParams.m_nOutput, outputs);return true;}

PoolTask::onDetectionResult

void PoolTask::onDetectionResult(TaskData &&task_data, std::vector<DetectionResult> &&dets){TaskResult res{true, std::move(task_data), std::move(dets)};list<TaskResult> tmp_res;{lock_guard<mutex> lg(m_mutexOutput);m_taskResults[res.m_taskData.m_taskId] = std::move(res);while (true){if (m_taskResults[m_outputIndex].m_isGet){auto &res_data = m_taskResults[m_outputIndex];tmp_res.emplace_back(std::move(res_data));res_data.m_isGet = false;++m_outputIndex;}else{break;}}}// 通过这个方式进行拼接，不在一个锁里面耗费时间if (!tmp_res.empty()){lock_guard<mutex> lg(m_mutexCb);m_cbResults.splice(m_cbResults.end(), tmp_res);}}