584 lines
13 KiB
Markdown
584 lines
13 KiB
Markdown
# YOLO 分割模型训练完整教程
|
||
|
||
## 📚 目录
|
||
|
||
1. [快速开始 - 训练搭电设备数据集](#快速开始)
|
||
2. [训练其他数据集](#训练其他数据集)
|
||
3. [详细配置说明](#详细配置说明)
|
||
4. [常见问题解决](#常见问题解决)
|
||
5. [训练技巧和优化](#训练技巧和优化)
|
||
|
||
---
|
||
|
||
## 🚀 快速开始 - 训练搭电设备数据集
|
||
|
||
### 步骤 1: 准备数据
|
||
|
||
确保你的数据集目录结构如下:
|
||
```
|
||
datasets/
|
||
└── 搭电设备/
|
||
├── data.yaml # 数据集配置文件
|
||
├── train/
|
||
│ ├── images/ # 训练图片(.jpg, .png等)
|
||
│ └── labels/ # 训练标注(.txt文件,多边形格式)
|
||
└── val/
|
||
├── images/ # 验证图片
|
||
└── labels/ # 验证标注
|
||
```
|
||
|
||
### 步骤 2: 检查 data.yaml 配置
|
||
|
||
打开 `datasets/搭电设备/data.yaml`,确保内容如下:
|
||
```yaml
|
||
train: train/images
|
||
val: val/images
|
||
nc: 2
|
||
names: ['JumperCable', 'NoJumperCable']
|
||
```
|
||
|
||
**说明:**
|
||
- `train`: 训练图片目录(相对于data.yaml的路径)
|
||
- `val`: 验证图片目录
|
||
- `nc`: 类别数量
|
||
- `names`: 类别名称列表(顺序对应类别ID:0, 1, 2...)
|
||
|
||
### 步骤 3: 验证标注格式(可选但推荐)
|
||
|
||
在训练前验证标注格式是否正确:
|
||
```bash
|
||
python verify_seg_labels.py datasets/搭电设备/train/labels datasets/搭电设备/train/images
|
||
```
|
||
|
||
如果看到 "✓ 格式正确",说明标注没问题。
|
||
|
||
### 步骤 4: 开始训练
|
||
|
||
**方法一:直接运行(使用默认配置)**
|
||
```bash
|
||
python train_segmentation.py
|
||
```
|
||
|
||
**方法二:修改脚本参数**
|
||
|
||
打开 `train_segmentation.py`,修改 `main()` 函数中的参数:
|
||
|
||
```python
|
||
def main():
|
||
train_segmentation_model(
|
||
dataset_path="datasets/搭电设备/data.yaml", # 数据集路径
|
||
model_name="yolov8n-seg.pt", # 模型选择
|
||
epochs=100, # 训练轮数
|
||
batch=8, # 批次大小
|
||
imgsz=640, # 图片尺寸
|
||
device="cpu" # 设备:cpu 或 cuda
|
||
)
|
||
```
|
||
|
||
### 步骤 5: 查看训练结果
|
||
|
||
训练完成后,模型保存在:
|
||
```
|
||
runs/segment/train/weights/
|
||
├── best.pt # 最佳模型(推荐使用)
|
||
└── last.pt # 最后一轮的模型
|
||
```
|
||
|
||
训练日志和可视化结果在:
|
||
```
|
||
runs/segment/train/
|
||
├── results.png # 训练曲线图
|
||
├── confusion_matrix.png # 混淆矩阵
|
||
└── ...
|
||
```
|
||
|
||
---
|
||
|
||
## 🎯 训练其他数据集
|
||
|
||
### 示例 1: 训练仪表盘数据集
|
||
|
||
#### 步骤 1: 检查数据集结构
|
||
```bash
|
||
# 查看数据集目录
|
||
ls datasets/仪表盘/
|
||
```
|
||
|
||
#### 步骤 2: 检查 data.yaml
|
||
打开 `datasets/仪表盘/data.yaml`,应该看到:
|
||
```yaml
|
||
train: train/images
|
||
val: val/images
|
||
nc: 2
|
||
names: ['Dashboard', 'NoDashboard']
|
||
```
|
||
|
||
#### 步骤 3: 创建训练脚本或修改现有脚本
|
||
|
||
**选项 A: 创建新脚本 `train_dashboard_seg.py`**
|
||
```python
|
||
from train_segmentation import train_segmentation_model
|
||
|
||
def main():
|
||
train_segmentation_model(
|
||
dataset_path="datasets/仪表盘/data.yaml",
|
||
model_name="yolov8n-seg.pt",
|
||
epochs=100,
|
||
batch=8,
|
||
imgsz=640,
|
||
device="cpu"
|
||
)
|
||
|
||
if __name__ == '__main__':
|
||
main()
|
||
```
|
||
|
||
**选项 B: 修改 `train_segmentation.py`**
|
||
```python
|
||
def main():
|
||
train_segmentation_model(
|
||
dataset_path="datasets/仪表盘/data.yaml", # 修改这里
|
||
model_name="yolov8n-seg.pt",
|
||
epochs=100,
|
||
batch=8,
|
||
imgsz=640,
|
||
device="cpu"
|
||
)
|
||
```
|
||
|
||
#### 步骤 4: 运行训练
|
||
```bash
|
||
python train_dashboard_seg.py
|
||
# 或
|
||
python train_segmentation.py
|
||
```
|
||
|
||
---
|
||
|
||
### 示例 2: 训练检测仪数据集
|
||
|
||
#### 步骤 1: 检查 data.yaml
|
||
```yaml
|
||
train: train/images
|
||
val: val/images
|
||
nc: 2
|
||
names: ['Detector', 'NoDetector']
|
||
```
|
||
|
||
#### 步骤 2: 创建训练脚本 `train_detector_seg.py`
|
||
```python
|
||
from train_segmentation import train_segmentation_model
|
||
|
||
def main():
|
||
train_segmentation_model(
|
||
dataset_path="datasets/检测仪/data.yaml",
|
||
model_name="yolov8n-seg.pt",
|
||
epochs=100,
|
||
batch=8,
|
||
imgsz=640,
|
||
device="cpu"
|
||
)
|
||
|
||
if __name__ == '__main__':
|
||
main()
|
||
```
|
||
|
||
#### 步骤 3: 运行
|
||
```bash
|
||
python train_detector_seg.py
|
||
```
|
||
|
||
---
|
||
|
||
### 示例 3: 训练新数据集(从头开始)
|
||
|
||
假设你要训练一个名为 "新设备" 的数据集:
|
||
|
||
#### 步骤 1: 创建数据集目录结构
|
||
```bash
|
||
mkdir -p datasets/新设备/train/images
|
||
mkdir -p datasets/新设备/train/labels
|
||
mkdir -p datasets/新设备/val/images
|
||
mkdir -p datasets/新设备/val/labels
|
||
```
|
||
|
||
#### 步骤 2: 准备数据
|
||
- 将训练图片放到 `datasets/新设备/train/images/`
|
||
- 将训练标注(.txt文件)放到 `datasets/新设备/train/labels/`
|
||
- 将验证图片放到 `datasets/新设备/val/images/`
|
||
- 将验证标注放到 `datasets/新设备/val/labels/`
|
||
|
||
**重要:** 标注文件名必须与图片文件名一致(仅扩展名不同)
|
||
- 图片:`image001.jpg`
|
||
- 标注:`image001.txt`
|
||
|
||
#### 步骤 3: 创建 data.yaml
|
||
创建文件 `datasets/新设备/data.yaml`:
|
||
```yaml
|
||
train: train/images
|
||
val: val/images
|
||
nc: 3
|
||
names: ['设备A', '设备B', '设备C']
|
||
```
|
||
|
||
**说明:**
|
||
- `nc: 3` 表示有3个类别
|
||
- `names` 列表中的顺序对应类别ID:
|
||
- '设备A' = 类别 0
|
||
- '设备B' = 类别 1
|
||
- '设备C' = 类别 2
|
||
|
||
#### 步骤 4: 验证标注
|
||
```bash
|
||
python verify_seg_labels.py datasets/新设备/train/labels datasets/新设备/train/images
|
||
```
|
||
|
||
#### 步骤 5: 创建训练脚本
|
||
创建 `train_new_device.py`:
|
||
```python
|
||
from train_segmentation import train_segmentation_model
|
||
|
||
def main():
|
||
train_segmentation_model(
|
||
dataset_path="datasets/新设备/data.yaml",
|
||
model_name="yolov8n-seg.pt",
|
||
epochs=100,
|
||
batch=8,
|
||
imgsz=640,
|
||
device="cpu"
|
||
)
|
||
|
||
if __name__ == '__main__':
|
||
main()
|
||
```
|
||
|
||
#### 步骤 6: 开始训练
|
||
```bash
|
||
python train_new_device.py
|
||
```
|
||
|
||
---
|
||
|
||
## ⚙️ 详细配置说明
|
||
|
||
### 1. 模型选择
|
||
|
||
YOLO 提供了多种模型大小,根据你的需求选择:
|
||
|
||
| 模型 | 速度 | 精度 | 参数量 | 适用场景 |
|
||
|------|------|------|--------|----------|
|
||
| yolov8n-seg.pt | 最快 | 较低 | 最少 | 实时检测、资源受限 |
|
||
| yolov8s-seg.pt | 快 | 中等 | 较少 | 平衡速度和精度 |
|
||
| yolov8m-seg.pt | 中等 | 较高 | 中等 | 一般应用 |
|
||
| yolov8l-seg.pt | 较慢 | 高 | 较多 | 高精度需求 |
|
||
| yolov8x-seg.pt | 最慢 | 最高 | 最多 | 最高精度需求 |
|
||
|
||
**修改方法:**
|
||
```python
|
||
model_name="yolov8s-seg.pt" # 改为你想要的模型
|
||
```
|
||
|
||
### 2. 训练轮数 (epochs)
|
||
|
||
- **小数据集(< 500张)**: 100-200 epochs
|
||
- **中等数据集(500-2000张)**: 50-100 epochs
|
||
- **大数据集(> 2000张)**: 30-50 epochs
|
||
|
||
**修改方法:**
|
||
```python
|
||
epochs=100 # 根据数据集大小调整
|
||
```
|
||
|
||
### 3. 批次大小 (batch)
|
||
|
||
根据你的GPU/CPU内存调整:
|
||
|
||
| 设备 | 推荐batch大小 |
|
||
|------|--------------|
|
||
| CPU | 4-8 |
|
||
| GPU 4GB | 8-16 |
|
||
| GPU 8GB | 16-32 |
|
||
| GPU 16GB+ | 32-64 |
|
||
|
||
**修改方法:**
|
||
```python
|
||
batch=8 # 如果内存不足,减小这个值;如果有GPU,可以增大
|
||
```
|
||
|
||
### 4. 图片尺寸 (imgsz)
|
||
|
||
- **640**: 标准尺寸,平衡速度和精度(推荐)
|
||
- **416**: 更快,但精度略低
|
||
- **832**: 更慢,但精度更高
|
||
|
||
**修改方法:**
|
||
```python
|
||
imgsz=640 # 可以改为 416, 512, 832 等
|
||
```
|
||
|
||
### 5. 设备选择 (device)
|
||
|
||
```python
|
||
device="cpu" # 使用CPU(较慢)
|
||
device="cuda" # 使用默认GPU
|
||
device="0" # 使用第1块GPU
|
||
device="1" # 使用第2块GPU
|
||
```
|
||
|
||
**检查GPU是否可用:**
|
||
```python
|
||
import torch
|
||
print(torch.cuda.is_available()) # True表示有GPU
|
||
```
|
||
|
||
---
|
||
|
||
## 🔧 常见问题解决
|
||
|
||
### 问题 1: 找不到模型文件
|
||
|
||
**错误信息:**
|
||
```
|
||
FileNotFoundError: yolov8n-seg.pt not found
|
||
```
|
||
|
||
**解决方法:**
|
||
- 首次运行会自动下载,确保网络连接正常
|
||
- 或手动下载后放到项目根目录
|
||
- 下载地址:https://github.com/ultralytics/assets/releases
|
||
|
||
### 问题 2: 内存不足 (Out of Memory)
|
||
|
||
**错误信息:**
|
||
```
|
||
RuntimeError: CUDA out of memory
|
||
```
|
||
|
||
**解决方法:**
|
||
1. 减小批次大小:
|
||
```python
|
||
batch=4 # 从8改为4
|
||
```
|
||
2. 减小图片尺寸:
|
||
```python
|
||
imgsz=416 # 从640改为416
|
||
```
|
||
3. 使用更小的模型:
|
||
```python
|
||
model_name="yolov8n-seg.pt" # 使用nano版本
|
||
```
|
||
|
||
### 问题 3: 标注文件格式错误
|
||
|
||
**错误信息:**
|
||
```
|
||
ValueError: invalid literal for int() with base 10
|
||
```
|
||
|
||
**解决方法:**
|
||
1. 运行验证脚本检查:
|
||
```bash
|
||
python verify_seg_labels.py datasets/你的数据集/train/labels
|
||
```
|
||
2. 确保标注格式为:`class_id x1 y1 x2 y2 x3 y3 ...`
|
||
3. 确保坐标值在 [0, 1] 范围内
|
||
|
||
### 问题 4: 训练损失不下降
|
||
|
||
**可能原因和解决方法:**
|
||
|
||
1. **学习率过高或过低**
|
||
- 使用默认学习率(通常不需要修改)
|
||
- 如果损失震荡,可以尝试减小学习率
|
||
|
||
2. **数据质量问题**
|
||
- 检查标注是否正确
|
||
- 确保训练集和验证集都有足够的样本
|
||
|
||
3. **模型太小**
|
||
- 尝试使用更大的模型(如 yolov8s-seg.pt)
|
||
|
||
4. **训练轮数不够**
|
||
- 增加 epochs 数量
|
||
|
||
### 问题 5: 验证集准确率很低
|
||
|
||
**解决方法:**
|
||
1. 检查验证集标注是否正确
|
||
2. 确保验证集和训练集的数据分布相似
|
||
3. 增加训练数据量
|
||
4. 使用数据增强(YOLO默认已启用)
|
||
|
||
### 问题 6: 训练速度很慢
|
||
|
||
**解决方法:**
|
||
1. **使用GPU**(最重要):
|
||
```python
|
||
device="cuda" # 改为GPU
|
||
```
|
||
2. 减小图片尺寸:
|
||
```python
|
||
imgsz=416 # 从640改为416
|
||
```
|
||
3. 使用更小的模型:
|
||
```python
|
||
model_name="yolov8n-seg.pt"
|
||
```
|
||
4. 启用缓存(如果内存足够):
|
||
```python
|
||
# 在train()中添加
|
||
cache=True
|
||
```
|
||
|
||
---
|
||
|
||
## 💡 训练技巧和优化
|
||
|
||
### 技巧 1: 使用预训练模型继续训练
|
||
|
||
如果训练中断,可以从上次保存的模型继续:
|
||
```python
|
||
# 加载之前训练的模型
|
||
model = YOLO("runs/segment/train/weights/last.pt")
|
||
|
||
# 继续训练
|
||
model.train(
|
||
data="datasets/搭电设备/data.yaml",
|
||
epochs=50, # 继续训练50轮
|
||
resume=True # 继续训练
|
||
)
|
||
```
|
||
|
||
### 技巧 2: 调整学习率
|
||
|
||
如果训练不稳定,可以调整学习率:
|
||
```python
|
||
results = model.train(
|
||
data="datasets/搭电设备/data.yaml",
|
||
epochs=100,
|
||
lr0=0.001, # 初始学习率(默认0.01)
|
||
lrf=0.1, # 最终学习率因子
|
||
...
|
||
)
|
||
```
|
||
|
||
### 技巧 3: 数据增强
|
||
|
||
YOLO默认启用数据增强,可以调整:
|
||
```python
|
||
results = model.train(
|
||
data="datasets/搭电设备/data.yaml",
|
||
epochs=100,
|
||
hsv_h=0.015, # 色调增强
|
||
hsv_s=0.7, # 饱和度增强
|
||
hsv_v=0.4, # 明度增强
|
||
degrees=10, # 旋转角度
|
||
translate=0.1, # 平移
|
||
scale=0.5, # 缩放
|
||
fliplr=0.5, # 水平翻转概率
|
||
...
|
||
)
|
||
```
|
||
|
||
### 技巧 4: 早停机制
|
||
|
||
如果验证损失不再下降,自动停止训练:
|
||
```python
|
||
results = model.train(
|
||
data="datasets/搭电设备/data.yaml",
|
||
epochs=100,
|
||
patience=50, # 50轮没有改善就停止
|
||
...
|
||
)
|
||
```
|
||
|
||
### 技巧 5: 多GPU训练
|
||
|
||
如果有多个GPU,可以使用:
|
||
```python
|
||
device=[0, 1] # 使用GPU 0和1
|
||
```
|
||
|
||
### 技巧 6: 监控训练过程
|
||
|
||
训练时会实时显示:
|
||
- 损失曲线
|
||
- 准确率
|
||
- mAP(平均精度)
|
||
|
||
训练完成后查看:
|
||
```bash
|
||
# 查看训练结果图片
|
||
open runs/segment/train/results.png
|
||
|
||
# 查看混淆矩阵
|
||
open runs/segment/train/confusion_matrix.png
|
||
```
|
||
|
||
---
|
||
|
||
## 📝 完整训练示例
|
||
|
||
### 示例:训练搭电设备数据集(完整流程)
|
||
|
||
```python
|
||
from train_segmentation import train_segmentation_model
|
||
|
||
def main():
|
||
# 步骤1: 验证标注(可选)
|
||
# python verify_seg_labels.py datasets/搭电设备/train/labels datasets/搭电设备/train/images
|
||
|
||
# 步骤2: 开始训练
|
||
train_segmentation_model(
|
||
dataset_path="datasets/搭电设备/data.yaml",
|
||
model_name="yolov8n-seg.pt", # 使用nano模型(速度快)
|
||
epochs=100, # 训练100轮
|
||
batch=8, # 批次大小8
|
||
imgsz=640, # 图片尺寸640
|
||
device="cuda" # 使用GPU(如果有)
|
||
)
|
||
|
||
# 步骤3: 训练完成后,模型保存在 runs/segment/train/weights/best.pt
|
||
# 可以使用这个模型进行推理
|
||
|
||
if __name__ == '__main__':
|
||
main()
|
||
```
|
||
|
||
---
|
||
|
||
## 🎓 总结
|
||
|
||
### 训练新数据集的通用步骤:
|
||
|
||
1. ✅ **准备数据** - 组织好目录结构
|
||
2. ✅ **创建 data.yaml** - 配置数据集信息
|
||
3. ✅ **验证标注** - 确保格式正确
|
||
4. ✅ **修改训练脚本** - 设置正确的数据集路径
|
||
5. ✅ **开始训练** - 运行训练脚本
|
||
6. ✅ **查看结果** - 检查训练曲线和模型
|
||
|
||
### 快速参考表:
|
||
|
||
| 要修改的内容 | 修改位置 | 示例 |
|
||
|------------|---------|------|
|
||
| 数据集路径 | `dataset_path` | `"datasets/仪表盘/data.yaml"` |
|
||
| 模型大小 | `model_name` | `"yolov8s-seg.pt"` |
|
||
| 训练轮数 | `epochs` | `150` |
|
||
| 批次大小 | `batch` | `16` |
|
||
| 图片尺寸 | `imgsz` | `832` |
|
||
| 设备类型 | `device` | `"cuda"` |
|
||
|
||
---
|
||
|
||
## 📞 需要帮助?
|
||
|
||
如果遇到问题:
|
||
1. 查看错误信息
|
||
2. 检查数据格式
|
||
3. 运行验证脚本
|
||
4. 参考本文档的"常见问题解决"部分
|
||
|
||
祝你训练顺利!🎉
|
||
|