# 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. 参考本文档的"常见问题解决"部分

祝你训练顺利！🎉