# YOLOv5 🚀 by Ultralytics, GPL-3.0 license """ Train a YOLOv5 model on a custom dataset. Models and datasets download automatically from the latest YOLOv5 release. Usage - Single-GPU training: $ python train.py --data coco128.yaml --weights yolov5s.pt --img 640 # from pretrained (recommended) $ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640 # from scratch Usage - Multi-GPU DDP training: $ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 train.py --data coco128.yaml --weights yolov5s.pt --img 640 --device 0,1,2,3 Models: https://github.com/ultralytics/yolov5/tree/master/models Datasets: https://github.com/ultralytics/yolov5/tree/master/data Tutorial: https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data """ import argparse import math import os import random import sys import time from copy import deepcopy from datetime import datetime from pathlib import Path import numpy as np import torch import torch
时间: 2025-06-25 08:01:13 浏览: 13
### YOLOv5 自定义数据集训练指南
#### 单GPU训练方法
YOLOv5 支持通过指定 `--device` 参数来选择单个 GPU 进行训练。以下是基于官方文档的命令示例[^1]:
```bash
python train.py --img 640 --batch 16 --epochs 50 --data custom_data.yaml --weights yolov5s.pt --device 0
```
在此命令中:
- `--img`: 设置输入图像大小,默认为 640。
- `--batch`: 批量大小,可根据显存调整。
- `--epochs`: 总共训练轮数。
- `--data`: 数据集配置文件路径,需提前准备好自定义数据集的 YAML 文件。
- `--weights`: 使用预训练权重文件作为初始模型参数。
- `--device`: 设定设备编号,此处设为 0 表示仅使用第一个 GPU。
#### 多GPU分布式训练(DDP)
对于多 GPU 场景,推荐采用分布式数据并行(DistributedDataParallel, DDP)模式进行训练[^2]。该模式能够显著提升大规模数据处理效率,并充分利用多个 GPU 的计算能力。启动脚本如下所示:
```bash
python -m torch.distributed.run --nproc_per_node=4 train.py --img 640 --batch 64 --epochs 50 --data custom_data.yaml --weights yolov5s.pt --device 0,1,2,3
```
此命令的关键部分解释如下:
- `-m torch.distributed.run`: 启动 PyTorch 分布式运行环境。
- `--nproc_per_node`: 指定每台机器使用的 GPU 数目,在这里设定为 4。
- `--device`: 列举所要分配的具体 GPU 编号列表。
需要注意的是,虽然支持旧版的 Multi-GPU DataParallel 方式,但由于其性能较差且易引发同步问题,因此不建议继续沿用[^2]。
#### 下载所需资源
为了顺利开展实验,需要获取两个主要组件——模型权重与 COCO 数据集标注样例[coco128.yaml][^3]。
##### 模型下载
访问 GitHub 官方仓库地址或者 Release 页面可以找到各种版本的预训练权重文件[yolov5s.pt][^3]。这些权重通常是在大型公开数据集上完成初步学习过程的结果,有助于加速收敛速度和提高泛化效果。
##### 数据集准备
COCO 是计算机视觉领域广泛认可的标准之一,而 coco128.yaml 提供了一个小型子集用于快速测试目的[^4]。可以通过克隆项目源码库自动拉取相关素材:
```bash
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt
wget https://raw.githubusercontent.com/ultralytics/yolov5/master/data/coco128.yaml
```
另外,针对个人定制化的场景,则应按照特定格式整理图片及其对应的标签信息到本地目录结构当中去[^3]。
---
###
阅读全文
相关推荐
















