🔥 项目简介

本项目基于 Ultralytics 框架深度改造，将其从目标检测框架扩展为统一的语义分割训练框架。核心设计目标不是支持单一模型，而是通过切换 YAML 配置文件，在同一套训练体系中快速训练、验证和对比多种不同结构的语义分割模型。
作者：AICurator
核心价值：

• 🚀 一键切换模型：无需修改训练代码，仅替换 YAML 即可训练不同网络
• 🧩 模型家族隔离：各模型实现独立封装，易于维护和扩展
• 📊 统一实验接口：相同的数据集、训练入口、指标输出和验证流程
• 🎯 保留 Ultralytics 风格：统一训练/验证接口、日志输出和结果保存
  完整工程下载地址：完整工程

🧠 支持的语义分割模型

本框架已集成以下主流语义分割模型家族，覆盖 CNN 和 Transformer 两大类架构：

模型家族	特点	骨干网络示例
UNet	经典编码器-解码器结构，带 Skip Connection	ResNet, VGG, MobileNet, EfficientNet, DenseNet
UNet++	嵌套跳跃连接，密集特征融合	ResNet34 等
DeepLabV3+	空洞空间金字塔池化，多尺度特征	ResNet, EfficientNet
DPT	Vision Transformer 骨干，密集预测	ViT-B/16-384 等
FPN	特征金字塔网络，多尺度融合	ResNet, EfficientNet
PSPNet	金字塔场景解析，全局上下文	ResNet 等
MAnet	多注意力机制	ResNet 等
PAN	路径聚合网络	ResNet 等
Linknet	轻量级编码器-解码器	ResNet 等
UPerNet	统一感知解析，Transformer 友好	ResNet 等
Segformer	轻量级 Transformer，高效语义分割	MiT-B0 ~ B5

💡 论文实验/模型对比友好：保持相同数据集和超参数，仅替换 YAML 即可完成多模型横向对比。

📁 工程目录结构

ultralytics-main/
├─ train.py                 # 统一训练入口
├─ val.py                   # 统一验证入口
├─ ultralytics/
│  ├─ cfg/
│  │  ├─ datasets/          # 数据集 YAML 配置
│  │  └─ models/            # 各模型家族 YAML 配置
│  │     ├─ unet/
│  │     ├─ unetplusplus/
│  │     ├─ deeplabv3plus/
│  │     ├─ dpt/
│  │     ├─ fpn/
│  │     ├─ pspnet/
│  │     ├─ manet/
│  │     ├─ pan/
│  │     ├─ linknet/
│  │     ├─ upernet/
│  │     └─ segformer/
│  └─ nn/
│     ├─ unet/              # UNet 模型实现
│     ├─ deeplabv3plus/     # DeepLabV3+ 模型实现
│     ├─ dpt/               # DPT 模型实现
│     ├─ fpn/               # FPN 模型实现
│     ├─ pspnet/            # PSPNet 模型实现
│     ├─ manet/             # MAnet 模型实现
│     ├─ pan/               # PAN 模型实现
│     ├─ linknet/           # Linknet 模型实现
│     ├─ upernet/           # UPerNet 模型实现
│     ├─ segformer/         # Segformer 模型实现
│     └─ losses/            # 语义分割损失函数集合
└─ runs/
   └─ semantic/             # 语义分割训练结果保存目录

⚡ 快速开始

环境安装

# 1. 创建虚拟环境（推荐 Python 3.9）
conda create -n py39 python=3.9 -y
conda activate py39

# 2. 安装 PyTorch（以 CUDA 12.4 为例）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# 3. 安装工程依赖
pip install -e .

# 4. 安装可选依赖
pip install timm segmentation-models-pytorch prettytable

数据集准备

数据集目录结构：

dataset/
└─ split_dataset/
   ├─ images/
   │  ├─ train/
   │  ├─ val/
   │  └─ test/
   └─ masks/
      ├─ train/
      ├─ val/
      └─ test/

要求：

• 图片和 mask 必须一一对应，文件名 stem 一致
• mask 为单通道灰度图，像素值表示类别 ID（0 ~ nc-1）
• 255 表示忽略区域

数据集 YAML 配置示例（ultralytics/cfg/datasets/my-semantic-seg.yaml）：

path: dataset/split_dataset

train: images/train
val: images/val
test: images/test

masks_train: masks/train
masks_val: masks/val
masks_test: masks/test

nc: 2

names:
  0: background
  1: spot

训练模型

方式一：使用 train.py（推荐）

修改 train.py 中的 yaml_path 即可切换模型：

# 训练 UNet + ResNet34
yaml_path = 'ultralytics/cfg/models/unet/unet-resnet34-expanded.yaml'

# 训练 DeepLabV3+ + EfficientNet-B0
yaml_path = 'ultralytics/cfg/models/deeplabv3plus/deeplabv3plus-efficientnet-b0.yaml'

# 训练 Segformer + MiT-B0
yaml_path = 'ultralytics/cfg/models/segformer/segformer-mit_b0-expanded.yaml'

运行：

python train.py

方式二：使用 Python API

from ultralytics import YOLO

# 加载模型配置
model = YOLO('ultralytics/cfg/models/unet/unet-resnet34-expanded.yaml')

# 开始训练
model.train(
    data='ultralytics/cfg/datasets/my-semantic-seg.yaml',
    task='semantic',
    imgsz=640,
    epochs=100,
    batch=4,
    loss='ce',           # 支持 ce/dice/focal/tversky/jaccard/lovasz/mcc/bce_dice/ce_dice/focal_dice
    augment=True,
    val=True,
    plots=True,
)

验证模型

python val.py

或修改 val.py 中的 model_path 指向训练好的权重：

model_path = 'runs/semantic/train39/weights/best.pt'

常用训练参数

参数	说明	示例
data	数据集 YAML 路径	'ultralytics/cfg/datasets/my-semantic-seg.yaml'
task	任务类型	'semantic'
imgsz	输入图像尺寸	640
epochs	训练轮数	100
batch	批大小	4
loss	损失函数	'ce' / 'dice' / 'focal' / 'ce_dice' 等
optimizer	优化器	'SGD' / 'Adam' / 'AdamW'
lr0	初始学习率	0.001
augment	是否开启数据增强	True / False
mosaic	Mosaic 增强概率	1.0
patience	早停耐心值	50

支持的损失函数

通过 loss 参数自由切换：

• ce — Cross Entropy
• dice — Dice Loss
• focal — Focal Loss
• tversky — Tversky Loss
• jaccard — Jaccard / IoU Loss
• lovasz — Lovász-Softmax
• mcc — Matthews Correlation Coefficient
• bce_dice — BCE + Dice 组合
• ce_dice — CE + Dice 组合
• focal_dice — Focal + Dice 组合

📊 训练结果

训练结果默认保存在 runs/semantic/ 目录下：

runs/semantic/train39/
├─ args.yaml          # 训练参数配置
├─ results.csv        # 每轮训练指标
├─ weights/
│  ├─ best.pt         # 最佳验证指标权重
│  └─ last.pt         # 最后一轮权重
└─ ...                # 训练曲线、可视化结果

验证指标说明

• 单类指标：IoU / Dice / Precision / Recall
• 全局指标：mIoU / fwIoU / mPA / Dice / Precision / PixAcc

验证结果同时保存至 paper_data.txt，便于论文数据整理。

💡 模型选择建议

场景	推荐模型
数据量不大，快速建立基线	UNet / UNet++ / DeepLabV3+
速度和效果平衡	FPN / PSPNet / Linknet
尝试 Transformer 架构	DPT / UPerNet / Segformer
论文多模型对比实验	保持数据集一致，仅替换 YAML

🔧 扩展新模型

本工程采用隔离式集成设计，新增模型家族时：

1. 在 ultralytics/cfg/models/ 下新增模型 YAML 目录
2. 在 ultralytics/nn/ 下新增模型实现目录
3. 遵循现有模型的注册和解析约定
4. 不直接污染已有模型家族的代码

这种结构的好处：

• 容易维护
• 容易继续扩展
• 出问题时更容易定位
• 不影响已经稳定的模型家族

⚠️ 注意事项

1. 路径格式：脚本中尽量使用正斜杠 /，避免 Windows 反斜杠转义问题
2. DPT 输入尺寸：某些 DPT 变体与 ViT 输入尺寸强相关（如 224、384），请确保 YAML 与 imgsz 匹配
3. 显存管理：不同模型显存需求差异大，显存紧张时优先减小 batch，再减小 imgsz
4. Batch 敏感性：部分带 BN 的模型在 batch=1 时不稳定，建议 batch >= 2
5. 环境一致性：训练、验证、导出建议在同一个 Conda 环境中完成

📜 许可证

本项目基于 Ultralytics 开源框架改造，遵循原项目的许可证条款。原始 Ultralytics 代码许可证详情请参见 LICENSE 文件。

🙏 致谢

本项目基于 Ultralytics YOLO 框架扩展开发，感谢 Ultralytics 团队提供的优秀基础架构。