PyTorch 语义分割

简介

本仓库是基于 PyTorch 的语义分割/场景解析实现。代码易于使用，可在多种数据集上进行训练和测试。代码库主要以 ResNet50/101/152 作为骨干网络，并可轻松适配其他基础分类模型。已实现的网络包括 PSPNet 和 PSANet，它们分别在 2016 年 ECCV16 ImageNet 场景解析挑战赛、2017 年 CVPR17 LSUN 语义分割挑战赛以及 2018 年 CVPR18 WAD 可行驶区域分割挑战赛中获得第一名。实验所用的数据集示例包括 ADE20K、PASCAL VOC 2012 和 Cityscapes。

更新

• 2020.05.15：master 分支，使用官方 nn.SyncBatchNorm，仅支持多进程训练，已在 PyTorch 1.4.0 上测试通过。
• 2019.05.29：1.0.0 分支，同时支持多线程训练（nn.DataParallel）和多进程训练（nn.parallel.DistributedDataParallel)（推荐），后者速度更快。使用来自 EncNet 和 apex 的 syncbn，已在 PyTorch 1.0.0 上测试通过。

使用说明

1. 亮点：

   • 使用官方 nn.SyncBatchNorm 进行快速多进程训练（nn.parallel.DistributedDataParallel)。
   • 通过精心设计的代码结构，实现了更好的复现效果。
   • 所有初始化模型、训练好的模型和预测结果均可在 这里 获取。

2. 需求：

   • 硬件：4–8 张 GPU（建议每张显存 ≥11G）。
   • 软件：PyTorch≥1.1.0、Python3、tensorboardX。

3. 克隆仓库：

   git clone https://github.com/hszhao/semseg.git
   

4. 训练：

   • 下载相关数据集，并按如下方式创建指向这些数据集的符号链接（您也可以直接修改 config 文件夹中的相关路径）：

     cd semseg
     mkdir -p dataset
     ln -s /path_to_ade20k_dataset dataset/ade20k
     

   • 下载 ImageNet 预训练 模型，并将其放入 initmodel 文件夹中用于权重初始化。请务必按照 FAQ.md 中详细说明的正确数据集格式操作。

   • 在配置文件中指定使用的 GPU，然后开始训练：

     sh tool/train.sh ade20k pspnet50
     

   • 如果您使用 SLURM 进行节点管理，请取消注释 train.sh 中的相关行，然后执行训练：

     sbatch tool/train.sh ade20k pspnet50
     

5. 测试：

   • 下载训练好的分割模型，并将其放置在配置文件中指定的文件夹中，或直接修改指定路径。

   • 进行完整测试（获取官方榜单上的性能指标）：

     sh tool/test.sh ade20k pspnet50
     

   • 快速演示：对单张图像进行推理：

     PYTHONPATH=./ python tool/demo.py --config=config/ade20k/ade20k_pspnet50.yaml --image=figure/demo/ADE_val_00001515.jpg TEST.scales '[1.0]'
     

6. 可视化：集成 tensorboardX 以获得更好的可视化效果。

   tensorboard --logdir=exp/ade20k
   

7. 其他：

   • 资源：GoogleDrive 链接 包含共享的模型、可视化预测结果和数据列表。
   • 模型：可访问 ImageNet 预训练模型和训练好的分割模型。请注意，我们的 ImageNet 预训练模型与原始 ResNet 实现相比，在前几层略有不同。
   • 预测：可查看多个模型的可视化预测结果。
   • 数据集：属性信息（names 和 colors）位于 dataset 文件夹中，部分示例列表也可供查阅。
   • 常见问题解答：FAQ.md。
   • 往期视频预测：高精度——PSPNet、PSANet；高效率——ICNet。

性能

说明：mIoU/mAcc/aAcc 分别表示平均 IoU、各类别的平均准确率和所有像素的准确率。ss 表示单尺度测试，ms 表示多尺度测试。训练时间是在配备 8 块 GeForce RTX 2080 Ti 显卡的服务器上测量的。不同数据集之间的通用参数如下：

• 训练参数：sync_bn(真)、scale_min(0.5)、scale_max(2.0)、rotate_min(-10)、rotate_max(10)、zoom_factor(8)、ignore_label(255)、aux_weight(0.4)、batch_size(16)、base_lr(1e-2)、power(0.9)、momentum(0.9)、weight_decay(1e-4)。
• 测试参数：ignore_label(255)、scales（单尺度：[1.0]，多尺度：[0.5 0.75 1.0 1.25 1.5 1.75])。

1. ADE20K： 训练参数：classes(150)、train_h(473/465-PSP/A)、train_w(473/465-PSP/A)、epochs(100)。 测试参数：classes(150)、test_h(473/465-PSP/A)、test_w(473/465-PSP/A)、base_size(512)。

   • 设置：在 train（20210 张图像）数据集上训练，在 val（2000 张图像）数据集上测试。

   网络	mIoU/mAcc/aAcc(ss)	mIoU/mAcc/pAcc(ms)	训练时间
   PSPNet50	0.4189/0.5227/0.8039.	0.4284/0.5266/0.8106.	14h
   PSANet50	0.4229/0.5307/0.8032.	0.4305/0.5312/0.8101.	14h
   PSPNet101	0.4310/0.5375/0.8107.	0.4415/0.5426/0.8172.	20h
   PSANet101	0.4337/0.5385/0.8102.	0.4414/0.5392/0.8170.	20h

2. PSACAL VOC 2012： 训练参数：classes(21)、train_h(473/465-PSP/A)、train_w(473/465-PSP/A)、epochs(50)。 测试参数：classes(21)、test_h(473/465-PSP/A)、test_w(473/465-PSP/A)、base_size(512)。

   • 设置：在 train_aug（10582 张图像）数据集上训练，在 val（1449 张图像）数据集上测试。

   网络	mIoU/mAcc/aAcc(ss)	mIoU/mAcc/pAcc(ms)	训练时间
   PSPNet50	0.7705/0.8513/0.9489.	0.7802/0.8580/0.9513.	3.3h
   PSANet50	0.7725/0.8569/0.9491.	0.7787/0.8606/0.9508.	3.3h
   PSPNet101	0.7907/0.8636/0.9534.	0.7963/0.8677/0.9550.	5h
   PSANet101	0.7870/0.8642/0.9528.	0.7966/0.8696/0.9549.	5h

3. Cityscapes： 训练参数：classes(19)、train_h(713/709-PSP/A)、train_w(713/709-PSP/A)、epochs(200)。 测试参数：classes(19)、test_h(713/709-PSP/A)、test_w(713/709-PSP/A)、base_size(2048)。

   • 设置：在 fine_train（2975 张图像）数据集上训练，在 fine_val（500 张图像）数据集上测试。

   网络	mIoU/mAcc/aAcc(ss)	mIoU/mAcc/pAcc(ms)	训练时间
   PSPNet50	0.7730/0.8431/0.9597.	0.7838/0.8486/0.9617.	7h
   PSANet50	0.7745/0.8461/0.9600.	0.7818/0.8487/0.9622.	7.5h
   PSPNet101	0.7863/0.8577/0.9614.	0.7929/0.8591/0.9638.	10h
   PSANet101	0.7842/0.8599/0.9621.	0.7940/0.8631/0.9644.	10.5h

引用

如果您觉得代码或训练好的模型有用，请考虑引用以下内容：

@misc{semseg2019,
  author={Zhao, Hengshuang},
  title={semseg},
  howpublished={\url{https://github.com/hszhao/semseg}},
  year={2019}
}
@inproceedings{zhao2017pspnet,
  title={Pyramid Scene Parsing Network},
  author={Zhao, Hengshuang and Shi, Jianping and Qi, Xiaojuan and Wang, Xiaogang and Jia, Jiaya},
  booktitle={CVPR},
  year={2017}
}
@inproceedings{zhao2018psanet,
  title={{PSANet}: Point-wise Spatial Attention Network for Scene Parsing},
  author={Zhao, Hengshuang and Zhang, Yi and Liu, Shu and Shi, Jianping and Loy, Chen Change and Lin, Dahua and Jia, Jiaya},
  booktitle={ECCV},
  year={2018}
}

问题

我们收集了一些常见问题解答（FAQ.md）。欢迎提交 Pull Request 或提出建议。联系方式：hengshuangzhao at gmail.com。

semseg 快速上手指南

semseg 是一个基于 PyTorch 的语义分割/场景解析开源工具库，实现了 PSPNet 和 PSANet 等经典网络。本指南帮助开发者快速完成环境配置、安装及基础训练与测试。

环境准备

系统要求

• 硬件：推荐 4-8 张 GPU（显存建议 >= 11GB），以支持高效的多进程训练。
• 操作系统：Linux (推荐)

软件依赖

• Python: Python 3
• PyTorch: 版本 >= 1.1.0 (推荐使用 1.4.0 及以上以利用官方 nn.SyncBatchNorm)
• 其他库: tensorboardX (用于可视化)

提示：国内用户可使用清华或阿里镜像源加速 PyTorch 及相关库的安装。

pip install torch torchvision tensorboardX -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆仓库 从 GitHub 拉取最新代码：

   git clone https://github.com/hszhao/semseg.git
   cd semseg
   

2. 准备数据集 下载所需数据集（如 ADE20K, PASCAL VOC, Cityscapes），并在 semseg 目录下创建软链接。以下以 ADE20K 为例：

   mkdir -p dataset
   ln -s /path_to_ade20k_dataset dataset/ade20k
   

   注：请替换 /path_to_ade20k_dataset 为你本地的实际数据路径。也可直接修改 config 文件夹中的路径配置。

3. 准备预训练模型 下载 ImageNet 预训练模型（用于权重初始化），将其放入 initmodel 文件夹中。

   • 下载地址：Google Drive Link
   • 注意：请确保数据格式符合项目要求（参考原项目文档细节）。

基本使用

1. 开始训练

配置好 GPU 信息后，运行训练脚本。以下示例使用 ADE20K 数据集和 PSPNet50 网络进行多进程训练：

sh tool/train.sh ade20k pspnet50

如果你使用 SLURM 集群管理系统，请先取消 train.sh 中相关行的注释，然后执行：

sbatch tool/train.sh ade20k pspnet50

2. 模型测试

下载训练好的分割模型并放置于配置文件指定的目录（或修改配置路径）。

• 全量测试（获取标准性能指标）：

  sh tool/test.sh ade20k pspnet50
  

• 单图快速演示： 对单张图片进行推理并查看效果：

  PYTHONPATH=./ python tool/demo.py --config=config/ade20k/ade20k_pspnet50.yaml --image=figure/demo/ADE_val_00001515.jpg TEST.scales '[1.0]'
  

3. 可视化监控

项目集成了 tensorboardX，可通过以下命令启动可视化界面，实时监控训练损失和精度：

tensorboard --logdir=exp/ade20k

然后在浏览器中访问显示的本地地址即可。