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