Yolov5s算法从训练到部署
发布网友
发布时间：2024-10-24 13:41
共1个回答
热心网友
时间：2024-11-09 15:37
本文详细介绍了使用PyTorch神经网络框架GPU版本，针对自定义数据集训练Yolov5s目标检测算法的全过程，并阐述了将训练结果转化为瑞芯微RK系列搭载的NPU加速单元可执行的rknn格式模型的步骤。
首先，搭建PyTorch GPU环境，通过Anaconda虚拟Python环境，确保安装了与显卡驱动版本兼容的PyTorch GPU版本。查看CUDA版本，确保安装的PyTorch兼容此版本。
创建Anaconda环境并安装PyTorch，使用命令行验证环境设置，确保使用的是GPU版本。在PyCharm中新建工程，选择已安装PyTorch的环境，通过验证脚本确保GPU支持。
以训练口罩检测算法为例，克隆Yolov5代码工程，安装依赖，制作包含戴/不戴口罩图像的数据集。使用OpenCV收集实时视频流，保存图片至指定文件夹，通过labelImg工具生成YOLO格式标签文件。
数据集划分成训练集与验证集，使用脚本随机划分，确保标签文件与图像集正确关联。修改工程文件配置，如预训练权重文件、数据文件、模型文件及超参数设置。
训练算法模型，若遇到问题，通过查看错误信息定位解决方案。训练完成后，使用tensorboard查看训练结果，测试算法模型在detect.py文件中的性能。
进行量化转换，将pt格式的模型转换为onnx格式，修改模型后处理部分。使用RKNN Toolkit将onnx模型转换为rknn格式，支持瑞芯微RK系列板子的uint8精度需求。通过Netron工具验证模型量化效果。
部署到瑞芯微RK系列板子，参考GitHub仓库中的npu推理例程，熟悉推理流程，逐步实现实时摄像头流推理等功能。
