TorchSeg

本项目旨在提供一个基于 PyTorch 的语义分割模型的快速、模块化参考实现。

亮点

• 模块化设计： 通过组合不同的组件，轻松构建自定义的语义分割模型。
• 分布式训练： 比多线程并行方法（nn.DataParallel）快 >60%，我们采用多进程并行方法。
• 多GPU训练与推理： 支持多种推理方式。
• 提供预训练模型，并实现了多种语义分割模型。

前置条件

• PyTorch 1.0
  • pip3 install torch torchvision
• Easydict
  • pip3 install easydict
• Apex
• Ninja
  • sudo apt-get install ninja-build
• tqdm
  • pip3 install tqdm

更新日志

v0.1.1 (2019年5月14日)

• 发布预训练模型及所有训练好的模型
• 添加适用于 ADE20K 数据集的 PSANet
• 增加对 CamVid 和 PASCAL-Context 数据集的支持
• 仅支持分布式训练方式

模型库

预训练模型

• ResNet18
• ResNet50
• ResNet101

支持的模型

• FCN
• DFN
• BiSeNet
• PSPNet
• PSANet

性能与基准测试

SS：单尺度；MSF：多尺度 + 翻转

PASCAL VOC 2012

方法	主干网络	训练集	评估集	单尺度 mIoU	多尺度 mIoU	模型
FCN-32s	R101_v1c	train_aug	val	71.26	-
DFN(论文)	R101_v1c	train_aug	val	79.67	80.6*
DFN(我们)	R101_v1c	train_aug	val	79.40	81.40	GoogleDrive

80.6^*: 该结果为论文中报道的进一步在 train 数据集上微调后的结果。

Cityscapes

非实时方法

方法	主干网络	OHEM	训练集	评估集	单尺度 mIoU	多尺度 mIoU	模型
DFN(论文)	R101_v1c	✗	train_fine	val	78.5	79.3
DFN(我们)	R101_v1c	✗	train_fine	val	79.09	80.41	GoogleDrive
DFN(我们)	R101_v1c	✓	train_fine	val	79.16	80.53	GoogleDrive
BiSeNet(论文)	R101_v1c	✓	train_fine	val	-	80.3
BiSeNet(我们)	R101_v1c	✓	train_fine	val	79.09	80.39	GoogleDrive
BiSeNet(论文)	R18	✓	train_fine	val	76.21	78.57
BiSeNet(我们)	R18	✓	train_fine	val	76.28	78.00	GoogleDrive
BiSeNet(论文)	X39	✓	train_fine	val	70.1	72
BiSeNet(我们)*	X39	✓	train_fine	val	70.32	72.06	GoogleDrive

实时方法

方法	主干网络	OHEM	训练集	评估集	平均 IoU	模型
BiSeNet(论文)	R18	✓	train_fine	val	74.8
BiSeNet(我们)	R18	✓	train_fine	val	74.83	GoogleDrive
BiSeNet(论文)	X39	✓	train_fine	val	69
BiSeNet(我们)*	X39	✓	train_fine	val	68.51	GoogleDrive

BiSeNet(我们)^*：由于我们未在 PyTorch 中对 Xception39 模型进行 ImageNet 预训练，本次实验是从零开始训练的。我们将在后续发布该模型的预训练权重及相关实验结果。

ADE

方法	主干网络	训练集	评估集	单尺度 mIoU	单尺度准确率	模型
PSPNet(论文)	R50_v1c	train	val	41.68	80.04
PSPNet(我们)	R50_v1c	train	val	41.65	79.74	GoogleDrive
PSPNet(论文)	R101_v1c	train	val	41.96	80.64
PSPNet(我们)	R101_v1c	train	val	42.89	80.55	GoogleDrive
PSANet(论文)	R50_v1c	train	val	41.92	80.17
PSANet(我们)*	R50_v1c	train	val	41.67	80.09	GoogleDrive
PSANet(论文)	R101_v1c	train	val	42.75	80.71
PSANet(我们)	R101_v1c	train	val	43.04	80.56	GoogleDrive

PSANet(我们)^*：论文中的原始 PSANet 使用过参数化的机制来构建注意力图，而我们仅预测与特征图大小相同的注意力图。其性能与原始版本几乎相当。

待办事项

• 提供全面的文档
• 支持更多语义分割模型
  • Deeplab v3 / Deeplab v3+
  • DenseASPP
  • EncNet
  • OCNet

训练

1. 创建数据集配置文件：train.txt、val.txt、test.txt
   文件结构：（以制表符分隔）

   图片路径   标注路径
   

2. 根据需求修改 config.py
3. 开始训练网络：

分布式训练

我们使用官方的 torch.distributed.launch 来启动多GPU训练。该工具函数会根据所需使用的GPU数量生成相应数量的Python进程，每个进程只使用一块GPU。

对于每次实验，只需运行以下脚本：

export NGPUS=8
python -m torch.distributed.launch --nproc_per_node=$NGPUS train.py

推理

在评估器中，我们基于多进程实现了多 GPU 推理。在推理阶段，程序会根据所需使用的 GPU 数量启动相应数量的 Python 进程，每个 Python 进程将在单个 GPU 上处理整个评估数据集的一个子集。

1. 在验证集上评估训练好的网络：

   python3 eval.py
   

2. 输入参数：

   使用方法：-e epoch_idx -d device_idx [--verbose ] 
   [--show_image] [--save_path Pred_Save_Path]
   

免责声明

本项目目前仍在积极开发中。因此，当前正常运行的功能在未来版本中可能会出现中断。不过，如果您在使用过程中遇到任何问题，请随时提交 issue。

引用

以下是 BibTeX 格式的引用参考。请注意，BibTeX 条目需要使用 url LaTeX 宏包。

如果本项目对您的研究有所帮助，请考虑在您的论文中引用该项目。

@misc{torchseg2019,
  author =       {Yu, Changqian},
  title =        {TorchSeg},
  howpublished = {\url{https://github.com/ycszen/TorchSeg}},
  year =         {2019}
}

如果 DFN 对您的研究有所帮助，请考虑在您的论文中引用该工作。

@inproceedings{yu2018dfn,
  title={Learning a Discriminative Feature Network for Semantic Segmentation},
  author={Yu, Changqian and Wang, Jingbo and Peng, Chao and Gao, Changxin and Yu, Gang and Sang, Nong},
  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2018}
}

如果 BiSeNet 对您的研究有所帮助，请考虑在您的论文中引用该工作。

@inproceedings{yu2018bisenet,
  title={Bisenet: Bilateral segmentation network for real-time semantic segmentation},
  author={Yu, Changqian and Wang, Jingbo and Peng, Chao and Gao, Changxin and Yu, Gang and Sang, Nong},
  booktitle={European Conference on Computer Vision},
  pages={334--349},
  year={2018},
  organization={Springer}
}

为什么取名“Furnace”？

“Furnace”意为“炼金炉”。我们每个人都是“炼金术士”，因此我希望每个人都能拥有一座优质的炼金炉，用来修炼“炼金术”。祝您成为一名出色的炼金术士。

TorchSeg 快速上手指南

TorchSeg 是一个基于 PyTorch 的语义分割模块化参考实现，支持分布式训练和多 GPU 推理，旨在提供快速、灵活的模型构建体验。

环境准备

系统要求

• PyTorch: 版本 1.0 及以上
• GPU: 支持 CUDA 的 NVIDIA 显卡（推荐用于分布式训练）
• 操作系统: Linux (Ubuntu 推荐)

前置依赖

请确保已安装以下依赖库：

pip3 install torch torchvision
pip3 install easydict
pip3 install tqdm
sudo apt-get install ninja-build

可选但推荐：安装 Apex 以启用混合精度训练加速。

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir ./

安装步骤

1. 克隆项目代码

   git clone https://github.com/ycszen/TorchSeg.git
   cd TorchSeg
   

2. 验证环境 确保上述依赖安装无误，项目无需额外的 setup.py 安装步骤，直接在项目目录下运行即可。

基本使用

1. 数据准备

在项目目录下创建数据集配置文件（如 train.txt, val.txt），格式如下（使用 Tab 分隔图片路径和标签路径）：

path-of-the-image   path-of-the-groundtruth

2. 配置修改

根据实际需求编辑 config.py 文件，设置数据集路径、模型骨干网络（Backbone）及训练参数。

3. 开始训练（分布式模式）

TorchSeg 默认使用多进程分布式训练，速度比 nn.DataParallel 快 60% 以上。假设使用 8 张 GPU：

export NGPUS=8
python -m torch.distributed.launch --nproc_per_node=$NGPUS train.py

注：若使用单卡，可将 NGPUS 设为 1。

4. 模型评估与推理

使用训练好的模型在验证集上进行评估或生成预测结果：

python3 eval.py -e <epoch_idx> -d <device_idx> --show_image --save_path ./predictions

• -e: 指定评估的 epoch 索引。
• -d: 指定使用的 GPU 设备索引。
• --show_image: 可视化预测结果。
• --save_path: 保存预测图像的路径。

5. 使用预训练模型

项目提供了在 PASCAL VOC、Cityscapes 和 ADE20K 数据集上预训练的模型（如 ResNet18/50/101 骨干的 FCN, DFN, BiSeNet, PSPNet, PSANet）。

• 下载链接见原文 "Model Zoo" 部分的 Google Drive 地址。
• 下载后将其路径配置到 config.py 或直接加载权重进行微调/推理。