UPSNet：统一的全景分割网络

简介

UPSNet 最初在一篇 CVPR 2019 口头报告 中被提出。

免责声明

本仓库已在 Python 3.6 和 PyTorch 0.4.1 环境下测试通过。模型训练使用了 16 块 GPU，并借助 horovod 进行分布式训练。它同样可以在 Python 2.7 / PyTorch 1.0 环境下，以及使用 4 块 GPU 的情况下运行。

许可证

© Uber, 2018–2019。根据 Uber 非商业许可协议授权。

引用 UPSNet

如果您在研究中发现 UPSNet 有所帮助，请考虑引用以下文献：

@inproceedings{xiong19upsnet,
    Author = {Yuwen Xiong, Renjie Liao, Hengshuang Zhao, Rui Hu, Min Bai, Ersin Yumer, Raquel Urtasun},
    Title = {UPSNet: A Unified Panoptic Segmentation Network},
    Conference = {CVPR},
    Year = {2019}
}

主要结果

COCO 2017 数据集（在 train-2017 上训练）

	测试集划分	PQ	SQ	RQ	PQTh	PQSt
UPSNet-50	val	42.5	78.0	52.4	48.5	33.4
UPSNet-101-DCN	test-dev	46.6	80.5	56.9	53.2	36.7

Cityscapes 数据集

	PQ	SQ	RQ	PQTh	PQSt
UPSNet-50	59.3	79.7	73.0	54.6	62.7
UPSNet-101-COCO (ms 测试)	61.8	81.3	74.8	57.6	64.8

软件要求

我们建议使用 Anaconda3，因为它已经包含了许多常用的软件包。

硬件要求

我们建议使用 4 到 16 块 GPU 来训练我们的模型，每块 GPU 至少需要 11 GB 显存。

安装

将此仓库克隆到 $UPSNet_ROOT 目录下。

运行 init.sh 脚本以构建必要的 C++/CUDA 模块并下载预训练模型。

对于 Cityscapes 数据集：

假设您已将 Cityscapes 数据集下载至 $CITYSCAPES_ROOT 目录，并生成了 TrainIds 标签图像，请在 $UPSNet_ROOT 下执行 ln -s $CITYSCAPES_ROOT data/cityscapes 创建软链接，然后运行 init_cityscapes.sh 脚本以准备 Cityscapes 数据集供 UPSNet 使用。

对于 COCO 数据集：

假设您已将 COCO 数据集下载至 $COCO_ROOT 目录，并且该目录下包含 annotations 和 images 文件夹，请在 $UPSNet_ROOT 下执行 ln -s $COCO_ROOT data/coco 创建软链接，然后运行 init_coco.sh 脚本以准备 COCO 数据集供 UPSNet 使用。

训练：

python upsnet/upsnet_end2end_train.py --cfg upsnet/experiments/$EXP.yaml

测试：

python upsnet/upsnet_end2end_test.py --cfg upsnet/experiments/$EXP.yaml

我们在 upsnet/experiments 文件夹中提供了若干配置文件，分别适用于 Cityscapes 和 COCO 数据集的 16 块或 4 块 GPU 训练场景。

模型权重

目前可复现论文中实验结果的模型权重已开放下载。请按照以下步骤使用这些权重：

运行 download_weights.sh 脚本以获取 Cityscapes 和 COCO 数据集的训练好的模型权重。

对于 Cityscapes 数据集：

python upsnet/upsnet_end2end_test.py --cfg upsnet/experiments/upsnet_resnet50_cityscapes_16gpu.yaml --weight_path ./model/upsnet_resnet_50_cityscapes_12000.pth

python upsnet/upsnet_end2end_test.py --cfg upsnet/experiments/upsnet_resnet101_cityscapes_w_coco_16gpu.yaml --weight_path ./model/upsnet_resnet_101_cityscapes_w_coco_3000.pth

对于 COCO 数据集：

python upsnet/upsnet_end2end_test.py --cfg upsnet/experiments/upsnet_resnet50_coco_16gpu.yaml --weight_path model/upsnet_resnet_50_coco_90000.pth

python upsnet/upsnet_end2end_test.py --cfg upsnet/experiments/upsnet_resnet101_dcn_coco_3x_16gpu.yaml --weight_path model/upsnet_resnet_101_dcn_coco_270000.pth

UPSNet 快速上手指南

UPSNet 是一个统一的全景分割网络（Unified Panoptic Segmentation Network），源自 CVPR 2019 口头报告论文。本指南帮助开发者快速搭建环境并运行模型。

环境准备

系统要求

• 操作系统: Linux (推荐)
• Python: 3.6 (已测试)，也支持 2.7
• PyTorch: 0.4.1 (已测试)，也支持 1.0
• GPU: 训练推荐 4~16 张显卡，显存至少 11GB
• 分布式训练: 使用 Horovod (多卡训练必需)

前置依赖

推荐使用 Anaconda3 管理环境，它已包含许多常用包。 请确保已安装 CUDA 和 cuDNN 以支持 GPU 加速。

注意: 由于该项目基于较旧版本的 PyTorch (0.4.1)，在新环境中安装可能遇到兼容性问题。建议创建独立的虚拟环境。

安装步骤

1. 克隆仓库 将代码克隆到本地目录（设为 $UPSNet_ROOT）：

   git clone <repository_url> $UPSNet_ROOT
   cd $UPSNet_ROOT
   

2. 初始化环境与编译模块 运行初始化脚本，该脚本将构建必要的 C++/CUDA 模块并下载预训练模型：

   bash init.sh
   

3. 准备数据集

   针对 Cityscapes 数据集: 假设已将数据集下载至 $CITYSCAPES_ROOT 且已生成 TrainIds 标签图像，执行以下命令建立软链接并准备数据：

   ln -s $CITYSCAPES_ROOT data/cityscapes
   bash init_cityscapes.sh
   

   针对 COCO 数据集: 假设已将数据集下载至 $COCO_ROOT 且目录下包含 annotations 和 images 文件夹，执行以下命令：

   ln -s $COCO_ROOT data/coco
   bash init_coco.sh
   

4. 下载模型权重 (可选) 如果需要复现论文结果，可运行以下脚本下载训练好的权重：

   bash download_weights.sh
   

基本使用

训练模型

使用配置文件启动端到端训练。以下示例使用 Cityscapes 配置（可根据需要替换为 COCO 配置）：

python upsnet/upsnet_end2end_train.py --cfg upsnet/experiments/$EXP.yaml

注：upsnet/experiments 文件夹下提供了针对不同数据集和显卡数量（4 卡或 16 卡）的配置文件。

测试模型

使用预训练权重进行推理测试。

Cityscapes 示例 (ResNet-50):

python upsnet/upsnet_end2end_test.py --cfg upsnet/experiments/upsnet_resnet50_cityscapes_16gpu.yaml --weight_path ./model/upsnet_resnet_50_cityscapes_12000.pth

COCO 示例 (ResNet-50):

python upsnet/upsnet_end2end_test.py --cfg upsnet/experiments/upsnet_resnet50_coco_16gpu.yaml --weight_path model/upsnet_resnet_50_coco_90000.pth

COCO 示例 (ResNet-101 DCN):

python upsnet/upsnet_end2end_test.py --cfg upsnet/experiments/upsnet_resnet101_dcn_coco_3x_16gpu.yaml --weight_path model/upsnet_resnet_101_dcn_coco_270000.pth