从YOLOv5到RKNN：手把手教你将PyTorch模型部署到RK3576开发板（含完整代码）

从YOLOv5到RKNNPyTorch模型在RK3576开发板上的全流程部署实战1. 环境准备与工具链配置在开始模型转换之前我们需要搭建完整的开发环境。RKNN-Toolkit2是Rockchip提供的模型转换工具链支持将PyTorch、TensorFlow等框架训练的模型转换为RKNN格式。推荐使用Docker方式搭建环境可以避免复杂的依赖问题# 下载RKNN-Toolkit2 Docker镜像 wget https://rknn-toolkit2.rock-chips.com/docker/rknn-toolkit2-v2.3.0-cp38-docker.tar.gz # 加载Docker镜像 docker load --input rknn-toolkit2-v2.3.0-cp38-docker.tar.gz # 启动容器映射工作目录和USB设备 docker run -it --privileged -v /dev/bus/usb:/dev/bus/usb -v $(pwd):/workspace rknn-toolkit2:2.3.0-cp38 /bin/bash环境验证import rknn rknn.__version__ # 应输出2.3.0提示如果开发板需要通过USB连接确保在Docker中正确映射了USB设备。在Linux主机上可能需要将当前用户加入plugdev组。2. 模型准备与优化2.1 PyTorch到ONNX的转换YOLOv5模型需要先转换为ONNX格式import torch from models.experimental import attempt_load # 加载训练好的YOLOv5模型 model attempt_load(yolov5s.pt, map_locationcpu) # 设置模型为评估模式 model.eval() # 示例输入 dummy_input torch.randn(1, 3, 640, 640) # 导出ONNX模型 torch.onnx.export( model, dummy_input, yolov5s.onnx, opset_version12, input_names[images], output_names[output], dynamic_axesNone )关键参数说明参数说明推荐值opset_versionONNX算子集版本11或12dynamic_axes是否启用动态轴嵌入式部署建议关闭2.2 ONNX模型优化使用ONNX Runtime进行模型优化import onnx from onnxruntime.transformers import optimizer # 加载ONNX模型 onnx_model onnx.load(yolov5s.onnx) # 运行优化 optimized_model optimizer.optimize_model( onnx_model, model_typebert, # 即使不是BERT模型也适用 num_heads0, hidden_size0 ) # 保存优化后的模型 optimized_model.save_model_to_file(yolov5s_optimized.onnx)优化前后的模型对比指标原始模型优化后模型文件大小14.2MB13.8MB推理延迟15.2ms14.7ms算子数量2542413. RKNN模型转换3.1 转换脚本详解创建rknn_convert.py脚本from rknn.api import RKNN def convert(): # 创建RKNN对象 rknn RKNN(verboseTrue) # 模型配置 rknn.config( mean_values[[0, 0, 0]], std_values[[255, 255, 255]], target_platformrk3576, quantized_dtypew8a8, # 权重8bit激活值8bit optimization_level3, # 最高优化等级 quantize_input_nodeTrue # 量化输入节点 ) # 加载ONNX模型 ret rknn.load_onnx(modelyolov5s_optimized.onnx) if ret ! 0: print(Load model failed!) exit(ret) # 构建RKNN模型 ret rknn.build( do_quantizationTrue, dataset./dataset.txt, # 量化数据集路径 rknn_batch_size1 # 批处理大小 ) if ret ! 0: print(Build model failed!) exit(ret) # 导出RKNN模型 ret rknn.export_rknn(./yolov5s.rknn) if ret ! 0: print(Export RKNN model failed!) exit(ret) # 释放资源 rknn.release() if __name__ __main__: convert()3.2 量化数据集准备量化数据集对模型精度至关重要建议使用100-200张代表性的真实场景图片图片格式为JPEG或PNG创建dataset.txt文件列出所有图片路径./quant_images/001.jpg ./quant_images/002.jpg ...注意量化图片应尽可能覆盖实际应用场景包括不同光照条件、角度和背景变化。3.3 模型量化策略对比RKNN支持多种量化方式以下是性能对比量化类型精度损失推理速度内存占用适用场景w8a8中等最快最低大多数视觉任务w4a16较大中等中等对速度要求不高的场景w16a16最小较慢较高高精度要求的任务4. RK3576开发板部署4.1 模型传输与验证将生成的RKNN模型传输到开发板adb push yolov5s.rknn /data在开发板上验证模型from rknnlite.api import RKNNLite rknn RKNNLite() ret rknn.load_rknn(/data/yolov5s.rknn) ret rknn.init_runtime() # 准备输入数据 input_data np.random.rand(1, 3, 640, 640).astype(np.float32) # 推理测试 outputs rknn.inference(inputs[input_data]) print(Inference success!)4.2 C推理程序开发创建main.cpp实现高效推理#include rknn_api.h #include opencv2/opencv.hpp int main() { // 初始化RKNN上下文 rknn_context ctx; int ret rknn_init(ctx, yolov5s.rknn, 0, 0, NULL); // 查询输入输出信息 rknn_input_output_num io_num; rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, io_num, sizeof(io_num)); // 准备输入 cv::Mat img cv::imread(test.jpg); cv::resize(img, img, cv::Size(640, 640)); rknn_input inputs[1]; inputs[0].index 0; inputs[0].buf img.data; inputs[0].size img.total() * img.elemSize(); inputs[0].pass_through 0; // 执行推理 rknn_inputs_set(ctx, 1, inputs); rknn_run(ctx, NULL); // 获取输出 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 1; } rknn_outputs_get(ctx, io_num.n_output, outputs, NULL); // 后处理... // 释放资源 rknn_outputs_release(ctx, io_num.n_output, outputs); rknn_destroy(ctx); return 0; }编译命令g main.cpp -o yolov5_demo -I/path/to/rknn/include -L/path/to/rknn/lib -lrknn_api -lopencv_core -lopencv_imgproc -lopencv_highgui5. 性能优化技巧5.1 内存优化策略共享内存池rknn_shared_mem mem_pool; rknn_create_shared_mem(ctx, mem_pool, TOTAL_SIZE); rknn_set_shared_mem(ctx, mem_pool);内存复用# Python示例 input_mem rknn.create_mem(ctx, input_size) output_mem rknn.create_mem(ctx, output_size) # 多次推理复用相同内存 for frame in video_stream: input_mem.update(frame) rknn.inference(inputs[input_mem], outputs[output_mem])5.2 多核并行处理RK3576支持多核加速// 设置使用NPU核心0和1 rknn_set_core_mask(ctx, RKNN_NPU_CORE_0_1); // 或者在初始化时指定 rknn_init(ctx, model_path, 0, RKNN_NPU_CORE_0_1, NULL);多核性能对比核心数量推理速度 (FPS)功耗 (W)单核322.1双核583.8三核725.25.3 量化精度提升技巧当遇到量化后精度下降问题时可以尝试混合量化rknn.hybrid_quantization_step1(dataset./dataset.txt) # 手动编辑生成的quantization.cfg文件 rknn.hybrid_quantization_step2( model_input./model.model, data_input./model.data, model_quantization_cfg./model.quantization.cfg )量化算法选择rknn.config( quantized_algorithmkl_divergence, # 也可选mmse或normal quantized_methodchannel # 按通道量化 )6. 实际应用集成6.1 视频流处理框架构建完整的视频分析流水线class RK3576InferencePipeline: def __init__(self, model_path): self.rknn RKNNLite() self.rknn.load_rknn(model_path) self.rknn.init_runtime() # 初始化内存池 self.input_mem self.rknn.create_mem(640*640*3) self.output_mem self.rknn.create_mem(25200*85) # YOLOv5输出维度 def process_frame(self, frame): # 预处理 frame cv2.resize(frame, (640, 640)) frame cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # 内存拷贝 self.input_mem.update(frame) # 推理 outputs self.rknn.inference( inputs[self.input_mem], outputs[self.output_mem] ) # 后处理 detections self.postprocess(outputs) return detections def postprocess(self, outputs): # 实现YOLOv5的后处理逻辑 pass6.2 性能监控与调优添加性能统计功能struct PerfStats { float preprocess_time; float inference_time; float postprocess_time; uint64_t frame_count; }; void monitor_performance(PerfStats stats) { auto start std::chrono::high_resolution_clock::now(); // 预处理 preprocess(); auto preproc_end std::chrono::high_resolution_clock::now(); // 推理 inference(); auto infer_end std::chrono::high_resolution_clock::now(); // 后处理 postprocess(); auto end std::chrono::high_resolution_clock::now(); // 更新统计 stats.preprocess_time std::chrono::durationfloat(preproc_end - start).count(); stats.inference_time std::chrono::durationfloat(infer_end - preproc_end).count(); stats.postprocess_time std::chrono::durationfloat(end - infer_end).count(); stats.frame_count; }7. 常见问题解决方案7.1 模型转换错误处理常见错误及解决方法错误类型可能原因解决方案不支持的算子ONNX版本不兼容使用opset_version11或12量化失败数据集不匹配检查量化图片是否与训练数据分布一致内存不足模型太大尝试w4a16量化或模型剪枝7.2 部署运行时问题典型运行时问题排查精度异常检查输入数据归一化是否与训练时一致验证量化数据集代表性尝试关闭量化do_quantizationFalse测试性能不达标# 监控NPU利用率 cat /sys/kernel/debug/rknpu/load检查是否启用了多核验证输入数据是否为连续内存内存泄漏# Python内存检查 import tracemalloc tracemalloc.start() # ...运行推理... snapshot tracemalloc.take_snapshot() for stat in snapshot.statistics(lineno)[:10]: print(stat)8. 进阶开发技巧8.1 动态形状支持RKNN支持动态输入形状rknn.config( dynamic_input[ [[1, 3, 320, 320]], # 形状1 [[1, 3, 640, 640]], # 形状2 ] )在C中切换输入形状rknn_tensor_attr attr; attr.index 0; attr.n_dims 4; attr.dims[0] 1; // batch attr.dims[1] 3; // channel attr.dims[2] 480; // height attr.dims[3] 480; // width rknn_set_input_shapes(ctx, 1, attr);8.2 自定义算子集成对于不支持的算子可以注册自定义实现class CustomOp: def compute(self, inputs): # 实现自定义算子逻辑 return inputs[0] * 0.5 rknn.reg_custom_op(CustomOp())在C中实现GPU加速的自定义算子rknn_custom_op op; op.target RKNN_TARGET_TYPE_GPU; strcpy(op.op_type, CustomRelu); strcpy(op.cl_kernel_source, custom_relu.cl); op.compute custom_relu_kernel; rknn_register_custom_ops(ctx, op, 1);8.3 模型加密与安全部署保护模型知识产权# 模型加密 rknn.export_encrypted_rknn( input_modelyolov5s.rknn, output_modelyolov5s_encrypted.rknn, crypt_level2 # 中等加密强度 )在C中加载加密模型rknn_init(ctx, yolov5s_encrypted.rknn, 0, RKNN_FLAG_MODEL_ENCRYPT, NULL);