通过正点原子的ATK-3568了解到了YOLOP,这里记录下训练及测试及在onnxruntime部署的过程。
步骤:训练->测试->转成onnx->onnxruntime部署测试
一、前言
YOLOP是华中科技大学研究团队在2021年开源的研究成果,其将目标检测/可行驶区域分割和车道线检测三大视觉任务同时放在一起处理,并且在Jetson TX2开发板子上能够达到23FPS。
论文标题:YOLOP: You Only Look Once for Panoptic Driving Perception
论文地址: https://arxiv.org/abs/2108.11250
官方代码: https://github.com/hustvl/YOLOP
网络结构
相比于其它论文,YOLOP的论文很容易阅读。YOLOP的核心亮点就是多任务学习,而各部分都是拿其它领域的成果进行缝合,其网络结构如下图所示:
三个子任务共用一个Backbone和Neck,然后分出来三个头来执行不同的任务。
二、环境搭建
1、平台
使用的是AutoDL平台,配置如下:
2、环境创建
创建虚拟环境
conda create -n yolop_env python=3.8 -y
激活环境
conda activate yolop_env
安装pytorch
根据官网Previous PyTorch Versions | PyTorch安装pytorch,这里选择1.7版本,其他版本测试时,在转换成onnx时出错,所以在这里直接安装1.7版本。
pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
下载源码,代码地址:
https://github.com/hustvl/YOLOP
安装yolop环境
pip install -r requirements.txt
测试
python tools/demo.py --source ./inference/images/0ace96c3-48481887.jpg
运行 python tools/demo.py --source ./inference/videos/1.mp4
测试图片:–source 图片路径 (或存放图片的文件夹路径)
测试视频:–source 视频路径 (或存放视频的文件夹路径)
相机实时:–source 0
默认用cpu测试效果,如果想改成gpu,修改demo.py 198行。
三、onnxruntime测试
先安装onnxruntime
pip install onnxruntime
测试
python test_onnx.py --weight yolop-640-640.onnx --img ./inference/images/adb4871d-4d063244.jpg
四、训练
训练前需要下载数据集,可以根据YOLOP里的readme.md里的地址下载,也可以下载下面的网盘地址,把数据解压到YOLOP目录下。
链接:https://pan.baidu.com/s/1DJirasFncnnf4auI-BxEHA
提取码:1234
在Train之前修改相关参数,修改文件./lib/config/default.py
1、修改第9行:
_C.WORKERS = 0
2、修改第51-56行,这是数据集的路径
_C.DATASET = CN(new_allowed=True)
_C.DATASET.DATAROOT = '/root/YOLOP-main/datasets/images' # the path of images folder
_C.DATASET.LABELROOT = '/root/YOLOP-main/datasets/det_annotations' # the path of det_annotations folder
_C.DATASET.MASKROOT = '/root/YOLOP-main/datasets/da_seg_annotations' # the path of da_seg_annotations folder
_C.DATASET.LANEROOT = '/root/YOLOP-main/datasets/ll_seg_annotations' # the path of ll_seg_annotations folder
3、修改ATCH_SIZE
修改96行,121行_C.TRAIN.BATCH_SIZE_PER_GPU 为合适大小
修改完后,就可以训练了,执行下面命令开始训练:
python tools/train.py
训练时间很久,根据default.py文件,epoch共240轮,大概8小时左右,在训练过程中有一点要注意,每训练一轮结束后,模型和一些图片都会保存下来,造成磁盘空间不够,有可能训练失败,我只是测试所以只训练了100轮。
五、转onnx模型及测试
转换成onnx模型前,先安装环境
pip install onnxruntime
pip install onnx
pip install onnx-simplifier
pip install onnxoptimizer
pip install numpy==1.20.0
安装完后,运行python export_onnx.py转成onnx模型
onnxruntime测试和上面一样,执行下面命令测试
python test_onnx.py --weight yolop-640-640.onnx --img ./inference/images/adb4871d-4d063244.jpg
onnx转换并测试成功,接下来就是把模型转成rknn并在rk3568上部署。
遇到的问题:
[ONNXRuntimeError] : 1 : FAIL : Node (Mul_948) Op (Mul) [ShapeInferenceError]
处理:把torch降成1.7