YOLOv5环境搭建、训练流程及tensorrt转换生成plan文件

一、背景

github官网yolov5，代码什么的从这个网站下

二、环境搭建

有两种环境搭建方式，一是用conda搭个虚拟环境，然后安装所有需要的库跟依赖等；二是用docker容器，下载英伟达的pytorch image镜像，里面已经装好了一部分需要的包，如torch等。建议用第二种。

1）.conda虚拟环境

conda create -n yolov5 python=3.8
conda activate yolov5

2）.利用pytorch镜像创建环境

到这个网址去下载NCG catalogpytorch的镜像，比如：

docker pull nvcr.io/nvidia/pytorch:22.02-py3

然后根据该镜像起一个容器container，下面是一个起容器的简单命令例子：

docker run --gpus all 
            -it --name=pytorch_2112  
            --volume=根据自己需求的地址名:/workspace  
            --restart=always  
	    -p 19000:19000 
	    -p 19001:19001 
	    -p 19002:19003 
            --ipc=host --shm-size=8g  
            --workdir=/workspace   
            4c14b66a4c09 
            /bin/bash

进入环境

git clone https://github.com/ultralytics/yolov5   #克隆
cd yolov5
pip install -r requirements.txt   #安装

三、数据集制作

用不同的标注手段得到的格式是不一样的，最终目的都是为了得到yolo要求的数据格式。
以下参考意义不大，只是个人记录
以shoes为例：
shoes_summary文件到了下面jpg里面
用的海康的软件进行标注，在label中生成一个大的json文件。
所以总体流程是：海康格式->voc->darknet。
总之就是先制作voc的dataset，再制作voc的label.
最后把voc变为darknet格式。（其实后面可以有时间再优化优化步骤）

1.create_voc_dataset.py修改

generate_train_val_test_txt("shoes_summary") 

2.hk_to_voc_label.py修改

hk_to_voc('./jpg/shoes_summary/Result/shoes_summary.json', "./voc_label/")

3.修改voc_to_darknet_dataset.py开头的内容 4.修改完毕后运行以下命令

cd generate_dataset/
python3 hk_to_voc_label.py
python3 create_voc_dataset.py
python3 voc_to_darknet_dataset.py

此时最终的darknet格式数据集在generate_dataset/shoes_summary文件夹下

四、训练流程 1.放数据集到指定位置

数据都放到固定文件夹下，在这里是yolov5同级文件夹datasets里面

2.根据数据集所在位置编写yaml文件

编写data/shoes.yaml,这里shoes可以自己指定名字，只要train的时候-data 指定就行,yaml示例如下：

train: ../datasets/shoes_summary/score/images/train/
val: ../datasets/shoes_summary/score/images/val/

nc: 3
names: ['people','boots','noboots']

3.训练

-data为刚刚指定的文件
-weight 为预训练模型权重
-cfg 为指定模型大小

训练命令示例：

python3 train.py --data data/shoes.yaml --cfg yolov5l.yaml --weights /workspace/yolov5/yolov5l.pt --batch-size 8  --epochs 500 

这里记得每生成最终的需要的pt文件，就要把之前生成的删一删，pt文件比较大，存多了容易出问题

五.生成engine/plan 1.生成wts文件

这个目的是将生成的pt文件转成wts文件。主要用gen_wts.py，这个脚本在tensorrtx的yolov5文件夹下，要先将gen_wts.py拷贝到yolov5文件夹下

model = torch.load('/nvidia-dev/yolov5-old/runs/train/exp74/weights/best.pt', map_location=device)['model'].float() # load to FP32

上面是要加载的pt文件

f = open('yolov5_shoes.wts', 'w')

这是生成的wts的名字，也可以不改，直接写成下面的命令(命令只是示例，要根据自己文件路径改)：

python3 /workspace/yolov5/gen_wts.py --weights /workspace/yolov5/runs/train/yolov5-shoes/weights/best.pt --output yolov5_shoes.wts

2.拷贝wts文件

将wts文件拷贝会tensorrtx/yolov5文件夹下

3.改yololayer.h文件

改yololayer.h文件的内容 长宽 num_classes 的数量等

4.生成engine/plan

mkdir build
cd build
cmake ..
make

生成plan的命令：

# ./yolov5 -s [.wts] [.engine] [s/m/l/x or c gd gw]  // serialize model to plan file
./yolov5 -s /workspace/tensorrtx/yolov5/yolov5_shoes.wts shoes.plan l

5.验证生成plan是否没问题

就用它来直接推理一些图片，看一下就知道了，samples是图片文件路径

./yolov5 -d shoes.plan samples

最终得到我们的plan文件，可以用做后续部署