AI项目十六：YOLOP 训练+测试+模型评估_yolop训练

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

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

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