YOLOv5训练coco128数据集流程

一、安装pytorch

1、创建虚拟环境

         使用Anaconda安装，建议先创建一个虚拟环境。启动Anaconda Prompt，在命令行输入：

         conda create -name torch python==3.7.11

2、安装pytorch

         进入pytorch官网，根据需要选择系统、安装包、语言和计算平台。

         通过conda安装CPU版pytorch：

         conda install pytorch torchvision torchaudio cpuonly -c pytorch

         安装其他版本pytorch：点击Previous versions of PyTorch，找到更早的版本，比如1.8.0。在Anaconda命令行输入：

         conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cpuonly -c pytorch

         如果安装报错，可尝试使用pip进行安装。

         pip install pytorch==1.8.0

3、部署YOLOv5

         YOLOv5开源地址：https://github.com/ultralytics/yolov5

         下载源码以后，使用pycharm打开train.py文件。在pycharm配置conda环境：依次点击File->settings->Project->Python Interpreter->Add Python Interpreter->Conda Environment->Existing environment选择Anaconda3->envs文件夹下对应虚拟环境名称文件夹下的python.exe文件。

4、修改配置参数

         在train.py中找到parse_opt，根据需要修改其中参数，一般只需要对前六个参数进行修改。

weights为训练参数保存名称，cfg设置模型配置文件，这里选yolov5s.yaml，在data中配置训练数据配置文件为coco128.yaml。epoch根据自己想要训练的轮数自行修改，batchsize大小根据电脑内存或显存大小修改。

5、模型训练

         运行train.py文件，使用coco128进行yolov5s训练，首次运行会自动下载coco128数据集，下载完成后自动开始模型训练。由于yolov5s模型和coco128数据集都比较小，使用cpu训练1个epoch也就一两分钟，每轮训练后更新模型验证精度。观察指标：精度'P'，召回率'R'，[email protected]和[email protected]:.95。

 

6、精度验证

         训练完成后，在runs->train->exp->weights文件夹下找到last.pt和best.pt文件，将best.pt文件拷贝至yolov5-master目录下，修改为yolov5s.pt。打开val.py运行，进行模型验证。

 

7、推理验证

         运行detect.py，进行推理验证。

推理结果：

 

8、模型导出

         模型验证和推理结果满足性能需求后，运行export.py可以将训练得到的pytorch权重pt文件导出为onnx格式，用于推理部署。导出过程如果出现报错，可尝试修改parse_opt中的配置参数'opset'，将default修改为12或其他onnx版本。