YOLOv5-PyTorch 使用教程

1. 项目目录结构及介绍

.
├── README.md         # 项目说明文档
├── requirements.txt   # 必要的依赖库列表
├── train.py          # 训练脚本
├── detect.py          # 检测脚本
├── models             # 模型定义目录
│   ├── yolov5s.pt     # 预训练模型权重
│   └── ...
├── data               # 数据集配置目录
│   ├── coco.yaml       # COCO数据集配置示例
│   └── voc.yaml        # PASCAL VOC数据集配置示例
└── utils               # 辅助工具函数目录
    ├── bbox           # 箱子操作相关工具
    ├── detect          # 检测相关的工具
    ├── dice            # 变形计算工具
    ├── plot             # 绘图工具
    └── train.py        # 训练相关的工具函数

目录结构中：

• README.md：项目的基本信息和简介。
• requirements.txt：列出所有必要的Python包，用于环境搭建。
• train.py：主要的训练脚本，负责加载数据、配置参数并执行模型训练。
• detect.py：检测脚本，可以用于对图片或视频进行目标检测。
• models：存储模型定义和预训练权重。
• data：包含数据集的配置文件，如COCO和PASCAL VOC的数据配置。
• utils：各种辅助工具，如处理边界框、训练过程中的函数等。

2. 项目启动文件介绍

2.1 train.py

train.py是项目的训练入口，它涵盖了以下关键功能：

• 加载配置文件（在"data"目录下的yaml文件）。
• 初始化网络模型（例如从models目录加载预训练权重）。
• 设置优化器、学习率调度器和其他训练参数。
• 加载数据集，包括预处理和数据增强。
• 开始训练循环，包括前向传播、损失计算、反向传播和优化步骤。
• 在训练过程中，定期保存模型权重并评估验证集上的性能。

可以通过指定参数来运行train.py，例如训练自己的数据集：

python train.py --cfg ./data/coco.yaml --weights yolov5s.pt --epochs 10

2.2 detect.py

detect.py用于图像或视频的目标检测。核心功能包括：

• 加载配置文件和模型权重。
• 接收图像路径，或者开启视频流进行实时检测。
• 对输入数据进行预处理。
• 运行模型进行预测。
• 将预测结果（边界框和类别）可视化到原始图像上。
• 输出检测结果。

运行图像检测的命令可能如下：

python detect.py --weights runs/train/exp/last.pt --img image.jpg

3. 项目的配置文件介绍

配置文件位于data目录下，例如coco.yaml和voc.yaml，这些文件定义了数据集的相关设置，包括：

• data: 包含train和val两个子字典，分别对应训练集和验证集的信息。
• batch_size: 训练时每批次的样本数。
• img_size: 输入图像的尺寸。
• conf_thres: 目标检测的置信度阈值。
• iou_thres: NMS非极大抑制的IoU阈值。
• classes: 定义了检测的目标类别。
• augmentations: 训练时的数据增强策略。

例如，在coco.yaml中，你会看到类似这样的配置：

data:
  train: coco.yaml
  val: coco.yaml
  test: coco.yaml
  img_size: 640
  batch_size: 16
  num_workers: 8
  rect: true
classes: 80

根据实际需求，你可以自定义配置文件以适应新的数据集或调整训练参数。