3.8 KiB
Yolov5s部署
实现步骤:
-
本地: best.pth
-
本地: best.pth -> best.onnx
-
本地 Linux: best.onnx -> best.rknn
-
RK3588 Linux: best.rknn 部署在开发板
1 获取 best.pth
下载对应版本的yolo5代码,训练得到best.pth
https://github.com/ultralytics/yolov5/tree/v6.0
2 best.pth 转化为 best.onnx
修改一段代码(注意,训练阶段不要修改)。将yolov5/models/yolo.py文件中class Detect(nn.Module)类的函数forword修改。
修改后将best.pt移动至yolov5/export.py同一文件夹下,在命令行调用以下命令:
python export.py --weights best.pt --img 640 --batch 1 --include onnx
或
python export.py --weights best.pt --include onnx
便可以得到成功转换的模型best.onnx。
(可能会报错说ONNX版本不支持,那么就将export.py中parse_opt()函数的--opset参数修改为12)
(也可能报错说不支持SiLU,把./models/common.py里的SiLU改成ReLU即可)
3 best.onnx 转化为 best.rknn
此处的转换工具只能在Linux系统上运行,在linux上下载(clone)转换工具rknn-toolkit2。并配置好相关环境,使用docker
或直接配置。
- Docker
根据rknn-toolkit2\docker\docker_file\ubuntu_18_04_cp36
目录下的Dockerfile
构建。
- 直接配置(我所采用的)
2.1 下载conda环境(有的话跳过)
wget -c https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
chmod 777 Miniconda3-latest-Linux-x86_64.sh
sh Miniconda3-latest-Linux-x86_64.sh
结束后重启命令行或执行source ~/.bashrc
。
2.2 创建名为rknntool的python3.8环境,并进入:
conda create -n rknntool python=3.8
conda activate rknntool
2.3 进入./doc目录
pip install -r requirements_cp38-1.5.0.txt
2.3进入./ packages目录
pip install rknn_toolkit2-1.5.0+1fa95b5c-cp38-cp38-linux_x86_64.whl
终端中运行python,输入from rknn.api import RKNN
,不报错即成功。
进入./examples/onnx/yolov5
,将上文得到的best.onnx文件和待预测图片(也可以没有)复制到此文件夹中,修改test.py文件。
修改完后执行python test.py
,这里的test.py构建了一个虚拟的NPU运行环境,模拟在RK3588上真实运行的情况。这时在当前文件夹./example/onnx/yolov5
中可以看到生成的best.rknn
模型和结果图片result.jpg。
4 RK3588 上部署 best.rknn
开发板环境还不完备,后面需要什么就安装什么,apt install xxx,例如git,gcc,g++等。
- 拉取官方demo:rockchip-linux/rknpu2到自己的目录下
git clone https://github.com/rockchip-linux/rknpu2
(注意此步骤不能windows上下载后传到开发板,要在开发板上拉取)
- 进入
/home/lhk/rknpu2/examples/rknn_yolov5_demo
目录。有需要的话则修改下图内容为自己的。 - 将转换后的best.rknn文件放在
model/RK3588
目录下,运行bash ./build-linux_RK3588.sh
,成功后生成install文件夹,进入install/rknn_yolov5_demo_Linux
。(需要gcc,g++环境,安装即可) - 在model目录下放入需要推理的图片test.jpg,运行
./rknn_yolov5_demo ./model/RK3588/best.rknn ./model/test.jpg
也可以支持h264、h265、rtsp视频流。例如h264视频流,指令:
./rknn_yolov5_video_demo model/RK3588/best.rknn model/test.h264 264