若该文为原创文章,转载请注明原文出处。

通过正点原子的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
 

如有侵权,或需要完整代码,请及时联系博主。

Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐