【YOLO V5】代码复现过程

最新推荐文章于 2024-07-24 10:35:13 发布

zy_destiny

最新推荐文章于 2024-07-24 10:35:13 发布

阅读量2.4k

点赞数 3

分类专栏：目标检测论文文章标签： YOLO 深度学习机器学习目标检测语义分割

本文链接：https://blog.csdn.net/qq_38308388/article/details/129090109

版权

目标检测论文专栏收录该内容

8 篇文章 0 订阅

订阅专栏

接上篇，讲到如何从mask转成YOLOv5训练需要的txt数据集格式，这篇就在此基础上进行模型训练预测和部署转换吧！

3.1OLOv5 PyTorch Hub 预测

数据集制作快速通道：YOLO格式数据集制作

1.环境准备

环境版本要求：python>=3.7

pytorch>=1.7

其他需要的安装包可以直接用requirements.txt下载安装

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

2.YOLO训练

2.1 数据集准备

rootpath
    --images
        --1.png
        --2.png
        ...

    --labels
        --1.txt
        --2.txt
        ...

images就不用多说了，里面就是存放原始训练影像，格式必须为png或者jpg；

labels里可以有两种格式.

如果是目标检测任务，那毋庸置疑就是标准yolo目标检测格式label了，一共5列，第1列存放每个目标的类别，第2-5列分别存放中心点坐标和宽度高度（center_x，center_y，wid ，hig）

注意：坐标点位置要归一化到（0-1）之间，也就是像素坐标分别除以宽度和高度。即

#center_x_pixl,center_y_pixl,width_pixl,higth_pixl表示中心点像素坐标和bbox的宽度和高度
#center_x，center_y,wid,hig表示归一化后的坐标，即txt里写的内容
#width，higth表示原始影像的宽度和高度

center_x = center_x_pixl / width
wid = width_pixl / width
center_y = center_y_pixl / higth
hig = higth_pixl / higth

如果是语义分割任务，与目标检测相同的是第1列仍然存放每个目标的类别，不同的是后面所有列分别存放每个边界点的坐标（center_x1，center_y1, center_x2，center_y2, ......, center_xn，center_yn），表示这些点组成了一个多边形目标。注意此时的边界点坐标也是归一化后的，归一化方法与目标检测方法一致。

2.2 data.yaml准备

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco128  # dataset root dir
train: images/train2017  # train images (relative to 'path') 128 images
val: images/train2017  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes (80 COCO classes)
names:
  0: person
  1: bicycle
  2: car
  ...
  77: teddy bear
  78: hair drier
  79: toothbrush

path：数据地址

train：训练数据文件夹名称

val：验证数据文件夹名称

names：类别名称对应，得从0开始

2.3 yolov5.yaml准备

以yolov5s为例，其实要改的就是一个numclass

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

# Parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 6, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 3, C3, [1024]],
   [-1, 1, SPPF, [1024, 5]],  # 9
  ]

# YOLOv5 v6.0 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, C3, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

nc改为自己的类别总数就可以了，其他不用动，0也算一类哦，比如类别是0，1，那nc就是2。

2.4 训练命令

单GPU训练

python train.py  --batch 64 --data coco.yaml --weights yolov5s.pt --device 0

多GPU训练

python -m torch.distributed.run --nproc_per_node 2 train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0,1

训练结果在run/train/目录下生成exp文件夹，若参数不变，会保留多个exp，不会覆盖。

3.YOLO预测

推理过程可以用两种方式，一种是用pytorch hub工具预测，一种是代码行预测

3.1OLOv5 PyTorch Hub 预测

import torch

# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s")  # or yolov5n - yolov5x6, custom

# Images
img = "https://ultralytics.com/images/zidane.jpg"  # or file, Path, PIL, OpenCV, numpy, list

# Inference
results = model(img)

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.

3.2代码预测

python detect.py --weights yolov5s.pt --source 0                               # webcam
                                               img.jpg                         # image
                                               vid.mp4                         # video
                                               screen                          # screenshot
                                               path/                           # directory
                                               list.txt                        # list of images
                                               list.streams                    # list of streams
                                               'path/*.jpg'                    # glob
                                               'https://youtu.be/Zgi9g1ksQHc'  # YouTube
                                               'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream

结果自动保存在runs/detect目录下

4.torch2onnx模型部署转换

把自己训练好的模型导出成onnx格式，方便部署


python export.py --weights yolov5s.pt --include torchscript onnx

注意：可以增加参数--half ，以达到半精度带出的目的，这样导出的文件比较小。

5.onnx模型推理预测

python detect.py --weights yolov5s.pt                 # PyTorch
                           yolov5s.torchscript        # TorchScript
                           yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                           yolov5s_openvino_model     # OpenVINO
                           yolov5s.engine             # TensorRT
                           yolov5s.mlmodel            # CoreML (macOS only)
                           yolov5s_saved_model        # TensorFlow SavedModel
                           yolov5s.pb                 # TensorFlow GraphDef
                           yolov5s.tflite             # TensorFlow Lite
                           yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
                           yolov5s_paddle_model       # PaddlePaddle

6.onnx模型运行验证

python val.py --weights yolov5s.pt                 # PyTorch
                        yolov5s.torchscript        # TorchScript
                        yolov5s.onnx               # ONNX Runtime or OpenCV DNN with --dnn
                        yolov5s_openvino_model     # OpenVINO
                        yolov5s.engine             # TensorRT
                        yolov5s.mlmodel            # CoreML (macOS Only)
                        yolov5s_saved_model        # TensorFlow SavedModel
                        yolov5s.pb                 # TensorFlow GraphDef
                        yolov5s.tflite             # TensorFlow Lite
                        yolov5s_edgetpu.tflite     # TensorFlow Edge TPU
                        yolov5s_paddle_model       # PaddlePaddle