在 PyTorch 中使用 MIT ADE20K 数据集进行语义分割

这是在 MIT ADE20K 场景解析数据集（http://sceneparsing.csail.mit.edu/）上实现的语义分割模型的 PyTorch 版本。

ADE20K 是由 MIT 计算机视觉团队发布的、最大的开源语义分割和场景解析数据集。请访问以下链接，查看我们基于 Caffe 和 Torch7 的数据集及其实现仓库： https://github.com/CSAILVision/sceneparsing

如果您只想体验我们的演示，请尝试此链接：http://scenesegmentation.csail.mit.edu。您可以上传自己的照片并进行场景解析！

您也可以在此处使用 Colab 笔记本进行试验，以尝试对图像进行分割的代码。

所有预训练模型均可在此处找到： http://sceneparsing.csail.mit.edu/model/pytorch

[从左至右：测试图像、真实标签、预测结果]

语义类别的颜色编码可在此处查看： https://docs.google.com/spreadsheets/d/1se8YEtb2detS7OuPE86fXGyD269pMycAWe2mtKUj2W8/edit?usp=sharing

更新

• 现已支持 HRNet 模型。
• 我们使用配置文件来存储原本位于参数解析器中的大部分选项。这些选项的定义详见 config/defaults.py。
• 我们在数据预处理方面遵循 PyTorch 的实践规范（RGB [0, 1] 范围，减去均值，除以标准差）。

亮点

PyTorch 中的同步批归一化

该模块在训练过程中会跨所有设备计算均值和标准差。我们通过实验发现，较大的批量大小对于分割任务至关重要。感谢 Jiayuan Mao 的慷慨贡献，详情请参阅 Synchronized-BatchNorm-PyTorch。

该实现易于使用，原因如下：

• 完全由 Python 编写，无需额外的 C++ 扩展库。
• 与 PyTorch 的原生实现完全兼容。具体而言，它使用无偏方差更新移动平均，并采用 sqrt(max(var, eps)) 而不是 sqrt(var + eps)。
• 效率较高，仅比非同步批归一化慢 20% 至 30%。

多 GPU 训练时输入的动态尺度

对于语义分割任务，保持图像的宽高比在训练过程中非常重要。因此，我们重新实现了 DataParallel 模块，使其能够以 Python 字典的形式将数据分发到多个 GPU 上，从而使每个 GPU 可以处理不同尺寸的图像。同时，数据加载器的工作方式也有所不同。

^{现在数据加载器的批次大小始终等于 GPU 的数量，每个元素都会被发送到一个 GPU。它也兼容多进程。需要注意的是，多进程数据加载器的文件索引存储在主进程中，这与我们希望每个工作进程维护自己文件列表的目标相矛盾。因此，我们采取了一个技巧：尽管主进程仍然会为 __getitem__ 函数提供索引，但我们直接忽略该请求，随机发送一个批次字典。此外，数据加载器派生的所有工作进程都具有相同的随机种子，如果不使用上述技巧，多个工作进程将会产生完全相同的数据。为此，我们在激活数据加载器的多进程功能之前，添加了一行代码来设置 numpy.random 的默认种子。}

当前最先进的模型

• PSPNet 是一种场景解析网络，通过金字塔池化模块（PPM）聚合全局特征表示。它是 ILSVRC'16 MIT 场景解析挑战赛的冠军模型。详细信息请参阅 https://arxiv.org/abs/1612.01105。
• UPerNet 是一种基于特征金字塔网络（FPN）和金字塔池化模块（PPM）的模型。它不需要膨胀卷积，而膨胀卷积是一种耗时且占用大量内存的操作。在不使用任何花哨技巧的情况下，它的性能可以媲美甚至超过 PSPNet，同时训练时间更短，所需的 GPU 内存也更少。详细信息请参阅 https://arxiv.org/abs/1807.10221。
• HRNet 是一种最近提出的模型，它在整个网络中始终保持高分辨率的特征表示，避免了传统的瓶颈设计。该模型在一系列像素级标注任务中取得了当前最优的性能。详细信息请参阅 https://arxiv.org/abs/1904.04514。

支持的模型

我们将模型分为编码器和解码器两部分，其中编码器通常直接基于分类网络修改而来，而解码器则由最后几层卷积和上采样组成。我们已在 config 文件夹中提供了一些预配置的模型。

编码器：

• MobileNetV2_dilated
• ResNet18/ResNet18_dilated
• ResNet50/ResNet50_dilated
• ResNet101/ResNet101_dilated
• HRNetV2 (W48)

解码器：

• C1（单个卷积模块）
• C1_deepsup（C1 加深监督技巧）
• PPM（金字塔池化模块，详情请参阅 PSPNet 论文）
• PPM_deepsup（PPM 加深监督技巧）
• UPerNet（金字塔池化结合 FPN 头部，详情请参阅 UperNet）

性能：

重要提示：我们仓库中的基础 ResNet 是自定义的（与 torchvision 中的版本不同）。基础模型会在需要时自动下载。

架构	多尺度测试	平均 IoU	像素准确率（%）	综合得分	推理速度（fps）
MobileNetV2dilated + C1_deepsup	否	34.84	75.75	54.07	17.2
	是	33.84	76.80	55.32	10.3
MobileNetV2dilated + PPM_deepsup	否	35.76	77.77	56.27	14.9
	是	36.28	78.26	57.27	6.7
ResNet18dilated + C1_deepsup	否	33.82	76.05	54.94	13.9
	是	35.34	77.41	56.38	5.8
ResNet18dilated + PPM_deepsup	否	38.00	78.64	58.32	11.7
	是	38.81	79.29	59.05	4.2
ResNet50dilated + PPM_deepsup	否	41.26	79.73	60.50	8.3
	是	42.14	80.13	61.14	2.6
ResNet101dilated + PPM_deepsup	否	42.19	80.59	61.39	6.8
	是	42.53	80.91	61.72	2.0
UperNet50	否	40.44	79.80	60.12	8.4
	是	41.55	80.23	60.89	2.9
UperNet101	否	42.00	80.79	61.40	7.8
	是	42.66	81.01	61.84	2.3
HRNetV2	否	42.03	80.77	61.40	5.8
	是	43.20	81.47	62.34	1.9

训练在配备 8 块 NVIDIA Pascal Titan Xp GPU（每块显存 12GB）的服务器上进行基准测试，推理速度则在单块 NVIDIA Pascal Titan Xp GPU 上进行基准测试，未启用可视化功能。

环境

代码在以下配置下开发：

• 硬件：训练至少 4 块 GPU，测试至少 1 块 GPU（相应设置 [--gpus GPUS]）
• 软件：Ubuntu 16.04.3 LTS，CUDA≥8.0，Python≥3.5，PyTorch≥0.4.0
• 依赖库：numpy、scipy、opencv、yacs、tqdm

快速入门：使用我们训练好的模型对一张图片进行测试

1. 以下是一个简单的示例，用于对单张图片进行推理：

chmod +x demo_test.sh
./demo_test.sh

该脚本会下载一个训练好的模型（ResNet50dilated + PPM_deepsup）和一张测试图片，运行测试脚本，并将预测的分割结果（.png）保存到工作目录中。

2. 若要对一张图片或一个包含多张图片的文件夹（$PATH_IMG）进行测试，只需执行以下命令：

python3 -u test.py --imgs $PATH_IMG --gpu $GPU --cfg $CFG

训练

1. 下载 ADE20K 场景解析数据集：

chmod +x download_ADE20K.sh
./download_ADE20K.sh

2. 通过选择要使用的 GPU（$GPUS) 和配置文件（$CFG) 来训练模型。训练过程中，默认情况下检查点会保存在 ckpt 文件夹中。

python3 train.py --gpus $GPUS --cfg $CFG 

• 要选择使用哪些 GPU，可以输入 --gpus 0-7，也可以输入 --gpus 0,2,4,6。

例如，您可以从我们提供的配置开始：

• 训练 MobileNetV2dilated + C1_deepsup

python3 train.py --gpus GPUS --cfg config/ade20k-mobilenetv2dilated-c1_deepsup.yaml

• 训练 ResNet50dilated + PPM_deepsup

python3 train.py --gpus GPUS --cfg config/ade20k-resnet50dilated-ppm_deepsup.yaml

• 训练 UPerNet101

python3 train.py --gpus GPUS --cfg config/ade20k-resnet101-upernet.yaml

3. 您还可以在命令行中覆盖选项，例如 python3 train.py TRAIN.num_epoch 10。

评估

1. 在验证集上评估训练好的模型。在参数中添加 VAL.visualize True 可以输出如预告片所示的可视化结果。

例如：

• 评估 MobileNetV2dilated + C1_deepsup

python3 eval_multipro.py --gpus GPUS --cfg config/ade20k-mobilenetv2dilated-c1_deepsup.yaml

• 评估 ResNet50dilated + PPM_deepsup

python3 eval_multipro.py --gpus GPUS --cfg config/ade20k-resnet50dilated-ppm_deepsup.yaml

• 评估 UPerNet101

python3 eval_multipro.py --gpus GPUS --cfg config/ade20k-resnet101-upernet.yaml

与其他项目的集成

可以通过 pip 安装此库，以便轻松地与其他代码库集成。

pip install git+https://github.com/CSAILVision/semantic-segmentation-pytorch.git@master

现在，此库可以方便地以编程方式使用。例如：

from mit_semseg.config import cfg
from mit_semseg.dataset import TestDataset
from mit_semseg.models import ModelBuilder, SegmentationModule

参考文献

如果您觉得代码或预训练模型有用，请引用以下论文：

通过 ADE20K 数据集实现场景的语义理解。B. Zhou, H. Zhao, X. Puig, T. Xiao, S. Fidler, A. Barriuso 和 A. Torralba。《国际计算机视觉期刊》（IJCV），2018 年。（https://arxiv.org/pdf/1608.05442.pdf）

@article{zhou2018semantic,
  title={Semantic understanding of scenes through the ade20k dataset},
  author={Zhou, Bolei and Zhao, Hang and Puig, Xavier and Xiao, Tete and Fidler, Sanja and Barriuso, Adela and Torralba, Antonio},
  journal={International Journal on Computer Vision},
  year={2018}
}

通过 ADE20K 数据集进行场景解析。B. Zhou, H. Zhao, X. Puig, S. Fidler, A. Barriuso 和 A. Torralba。计算机视觉与模式识别会议（CVPR），2017 年。（http://people.csail.mit.edu/bzhou/publication/scene-parse-camera-ready.pdf）

@inproceedings{zhou2017scene,
    title={Scene Parsing through ADE20K Dataset},
    author={Zhou, Bolei and Zhao, Hang and Puig, Xavier and Fidler, Sanja and Barriuso, Adela and Torralba, Antonio},
    booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
    year={2017}
}

semantic-segmentation-pytorch 快速上手指南

本指南基于 MIT ADE20K 数据集的 PyTorch 语义分割实现，帮助开发者快速部署和测试预训练模型。

1. 环境准备

系统要求

• 操作系统: Ubuntu 16.04.3 LTS 或更高版本
• 硬件:
  • 训练：建议 >= 4 块 GPU
  • 推理/测试：>= 1 块 GPU
• 核心软件:
  • CUDA >= 8.0
  • Python >= 3.5
  • PyTorch >= 0.4.0

依赖安装

确保已安装以下 Python 库：

pip install numpy scipy opencv-python yacs tqdm

提示：国内用户建议使用清华或阿里镜像源加速安装，例如： pip install -i https://pypi.tuna.tsinghua.edu.cn/simple numpy scipy opencv-python yacs tqdm

2. 安装步骤

你可以选择直接克隆仓库使用，或通过 pip 集成到其他项目中。

方式一：克隆仓库（推荐用于完整功能）

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

方式二：通过 pip 安装（用于项目集成）

pip install git+https://github.com/CSAILVision/semantic-segmentation-pytorch.git@master

注意：基础 ResNet 模型会在首次运行时自动下载，无需手动干预。

3. 基本使用

场景 A：快速测试单张图片（最简单示例）

项目提供了一个脚本，自动下载预训练模型（ResNet50dilated + PPM_deepsup）和测试图片并运行推理。

1. 赋予执行权限并运行演示脚本：

chmod +x demo_test.sh
./demo_test.sh

运行结束后，预测的分割结果（.png 格式）将保存在当前工作目录。

场景 B：自定义图片推理

如果你有自己的图片或文件夹需要测试，请使用以下命令：

python3 -u test.py --imgs $PATH_IMG --gpu $GPU --cfg $CFG

• $PATH_IMG: 图片路径或包含图片的文件夹路径。
• $GPU: 指定使用的 GPU ID（例如 0）。
• $CFG: 配置文件路径（例如 config/ade20k-resnet50dilated-ppm_deepsup.yaml）。

场景 C：代码集成示例

若已通过 pip 安装，可在 Python 代码中直接调用：

from mit_semseg.config import cfg
from mit_semseg.dataset import TestDataset
from mit_semseg.models import ModelBuilder, SegmentationModule

# 此处可继续编写加载模型和推理的逻辑

---

注：如需训练新模型或评估验证集性能，请参考官方 README 中的 Training 和 Evaluation 章节。