使用YOLOv5训练自定义数据集
在 Ubuntu20.04 系统上使用4张GPU卡基于容器训练模型。
准备数据
project
├── data.yaml #数据集配置文件
├── models #网络模型(可以使用下面的脚本自动生成)
│ ├── yolov5s.yaml #Small
│ ├── yolov5m.yaml #Medium
│ ├── yolov5l.yaml #Large
│ └── yolov5x.yaml #XLarge
├── images #图片
│ ├── train #训练集
│ │ ├── 000001.jpg
│ │ ├── 000002.jpg
│ │ └── 000003.jpg
│ └── val #验证集
│ ├── 000010.jpg
│ └── 000011.jpg
├── labels #YOLO格式的标注
│ ├── train #训练集
│ │ ├── 000001.txt
│ │ ├── 000002.txt
│ │ └── 000003.txt
│ └── val #验证集
│ ├── 000010.txt
│ └── 000011.txt
├── test #测试图片
└── output #推理后的标注图片
构建环境
- 拉取镜像
$ docker pull ultralytics/yolov5:latest
- 运行容器
$ docker run --ipc=host --runtime=nvidia -it --name project_name-yolov5 \
-v project_dir:/usr/src/app/project ultralytics/yolov5:latest