PyTorch 上的语义分割

英语 | 简体中文

该项目旨在提供一个简洁、易用且可修改的参考实现，用于在 PyTorch 中构建语义分割模型。

安装

# semantic-segmentation-pytorch 的依赖
pip install ninja tqdm

# 按照 https://pytorch.org/get-started/locally/ 中的说明安装 PyTorch
conda install pytorch torchvision -c pytorch

# 安装 PyTorch 语义分割库
git clone https://github.com/Tramac/awesome-semantic-segmentation-pytorch.git

使用方法

训练

---

• 单 GPU 训练

# 例如，训练 fcn32_vgg16_pascal_voc:
python train.py --model fcn32s --backbone vgg16 --dataset pascal_voc --lr 0.0001 --epochs 50

• 多 GPU 训练

# 例如，使用 4 个 GPU 训练 fcn32_vgg16_pascal_voc:
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS train.py --model fcn32s --backbone vgg16 --dataset pascal_voc --lr 0.0001 --epochs 50

评估

---

• 单 GPU 评估

# 例如，评估 fcn32_vgg16_pascal_voc
python eval.py --model fcn32s --backbone vgg16 --dataset pascal_voc

• 多 GPU 评估

# 例如，使用 4 个 GPU 评估 fcn32_vgg16_pascal_voc:
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS eval.py --model fcn32s --backbone vgg16 --dataset pascal_voc

演示

cd ./scripts
# 对于新用户：
python demo.py --model fcn32s_vgg16_voc --input-pic ../tests/test_img.jpg
# 您需要自行添加 'test.jpg'
python demo.py --model fcn32s_vgg16_voc --input-pic ../datasets/test.jpg

.{SEG_ROOT}
├── scripts
│   ├── demo.py
│   ├── eval.py
│   └── train.py

支持

模型

• FCN
• ENet
• PSPNet
• ICNet
• DeepLabv3
• DeepLabv3+
• DenseASPP
• EncNet
• BiSeNet
• PSANet
• DANet
• OCNet
• CGNet
• ESPNetv2
• DUNet(DUpsampling)
• FastFCN(JPU)
• LEDNet
• Fast-SCNN
• LightSeg
• DFANet

DETAILS 包含模型和骨干网络的详细信息。

.{SEG_ROOT}
├── core
│   ├── models
│   │   ├── bisenet.py
│   │   ├── danet.py
│   │   ├── deeplabv3.py
│   │   ├── deeplabv3+.py
│   │   ├── denseaspp.py
│   │   ├── dunet.py
│   │   ├── encnet.py
│   │   ├── fcn.py
│   │   ├── pspnet.py
│   │   ├── icnet.py
│   │   ├── enet.py
│   │   ├── ocnet.py
│   │   ├── psanet.py
│   │   ├── cgnet.py
│   │   ├── espnet.py
│   │   ├── lednet.py
│   │   ├── dfanet.py
│   │   ├── ......

数据集

您可以运行脚本来下载数据集，例如：

cd ./core/data/downloader
python ade20k.py --download-dir ../datasets/ade

数据集	训练集	验证集	测试集
VOC2012	1464	1449	✘
VOCAug	11355	2857	✘
ADK20K	20210	2000	✘
Cityscapes	2975	500	✘
COCO
SBU-shadow	4085	638	✘
LIP(Look into Person)	30462	10000	10000

.{SEG_ROOT}
├── core
│   ├── data
│   │   ├── dataloader
│   │   │   ├── ade.py
│   │   │   ├── cityscapes.py
│   │   │   ├── mscoco.py
│   │   │   ├── pascal_aug.py
│   │   │   ├── pascal_voc.py
│   │   │   ├── sbu_shadow.py
│   │   └── downloader
│   │       ├── ade20k.py
│   │       ├── cityscapes.py
│   │       ├── mscoco.py
│   │       ├── pascal_voc.py
│   │       └── sbu_shadow.py

结果

• PASCAL VOC 2012

方法	骨干网络	训练集	验证集	裁剪尺寸	epoch数	JPU	平均 IoU	像素准确率
FCN32s	vgg16	train	val	480	60	✘	47.50	85.39
FCN16s	vgg16	train	val	480	60	✘	49.16	85.98
FCN8s	vgg16	train	val	480	60	✘	48.87	85.02
FCN32s	resnet50	train	val	480	50	✘	54.60	88.57
PSPNet	resnet50	train	val	480	60	✘	63.44	89.78
DeepLabv3	resnet50	train	val	480	60	✘	60.15	88.36

注：lr=1e-4, batch_size=4, epochs=80。

过拟合测试

详情请参见 TEST。

.{SEG_ROOT}
├── tests
│   └── test_model.py

待办事项

• 添加训练脚本
• 移除 syncbn
• 进行训练和评估
• 测试分布式训练
• 修复 syncbn（为什么使用 SyncBN？）
• 添加分布式功能（如何进行 DIST？）

参考资料

• PyTorch-Encoding
• maskrcnn-benchmark
• gloun-cv
• imagenet

awesome-semantic-segmentation-pytorch 快速上手指南

本项目旨在提供一套简洁、易用且可修改的 PyTorch 语义分割模型参考实现，支持多种主流模型（如 FCN, PSPNet, DeepLabv3+ 等）及数据集。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐) 或 Windows
• Python: 2.x 或 3.x
• PyTorch: 1.1 及以上版本
• GPU: 推荐使用 NVIDIA GPU 以加速训练和推理

前置依赖:

• ninja: 用于加速编译
• tqdm: 用于显示进度条

安装步骤

1. 安装基础依赖

首先安装项目所需的 Python 依赖包：

pip install ninja tqdm

2. 安装 PyTorch

请访问 PyTorch 官网 获取适合您环境的安装命令。若使用 Conda，典型命令如下：

conda install pytorch torchvision -c pytorch

提示: 国内用户可使用清华或中科大镜像源加速下载，例如：

conda install pytorch torchvision -c pytorch --channel https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/

3. 克隆项目代码

将项目代码克隆到本地：

git clone https://github.com/Tramac/awesome-semantic-segmentation-pytorch.git
cd awesome-semantic-segmentation-pytorch

基本使用

1. 训练模型 (单卡示例)

以下命令演示如何使用 VGG16 作为骨干网络，在 Pascal VOC 数据集上训练 FCN32s 模型：

python train.py --model fcn32s --backbone vgg16 --dataset pascal_voc --lr 0.0001 --epochs 50

2. 评估模型

训练完成后，使用以下命令评估模型性能：

python eval.py --model fcn32s --backbone vgg16 --dataset pascal_voc

3. 推理演示 (Demo)

您可以使用预训练模型对单张图片进行语义分割预测。首先进入脚本目录，然后运行：

cd ./scripts
# 请将 ../tests/test_img.jpg 替换为您自己的图片路径
python demo.py --model fcn32s_vgg16_voc --input-pic ../tests/test_img.jpg

4. 数据集下载 (可选)

项目提供了脚本来自动下载常用数据集。以下以 ADE20K 为例：

cd ./core/data/downloader
python ade20k.py --download-dir ../datasets/ade

支持的数据集包括：Pascal VOC, ADE20K, Cityscapes, COCO, SBU-shadow, LIP 等。