pytorch-deeplab-resnet

本仓库包含用于我发表在 ACM MM 2017 上的论文《SketchParse: 基于多任务层次化深度网络实现对草图的丰富描述》中的 deepLab-ResNet 架构的代码。该架构在多个尺度（1倍、0.75倍、0.5倍）上对输入图像计算损失，这3个尺度分别独立计算损失。此外，在合并3个尺度的输出得分图之后，还会再计算一个损失。最终将这4个损失相加得到总损失。

更新

2017年6月24日

• 现在，3个尺度（1倍、0.75倍、0.5倍）的权重与 Caffe 实现一样是共享的。此前，每个尺度都有独立的权重。做出这一改动后，结果几乎相同（更多细节见结果部分）。不过，训练后的 .pth 模型文件大小显著减小。在 GPU 上占用的内存（11.9 GB）以及训练所需时间（约3.5小时）与之前保持一致。相应 .pth 文件的下载链接已更新。
• 可以使用自定义数据通过 train.py 训练 pytorch-deeplab-resnet，并为此在 train.py 和 evalpyt.py 中添加了 --NoLabels 标志（表示训练数据中的标签总数）。请注意，真值图像中的标签应使用连续数值（从0开始）。例如，如果有7种不同的标签，则每张真值图像必须包含这些标签：0,1,2,3,...6（即 no_labels-1）。

使用说明

请注意，本仓库仅在 Python 2.7 环境下进行了测试。

训练

步骤1：将 init.caffemodel 转换为 .pth 文件：init.caffemodel 包含 MS COCO 预训练的权重。我们使用这些权重作为除最后一层之外的所有层的初始化参数。对于最后一层，则采用标准差为0.01的随机高斯分布进行初始化。 要将 init.caffemodel 转换为 .pth 文件，请运行（或在此处下载转换好的 .pth 文件 链接）

python init_net_surgery.py

运行 init_net_surgery.py 需要 deeplab v2 caffe 和 PyTorch（Python 2.7）。

步骤2：现在我们已经有了初始化权重，可以通过运行以下命令来训练 deeplab-resnet：

python train.py

要查看每个命令行参数的说明，请运行：

python train.py -h

运行 train.py 需要 PyTorch（Python 2.7）。

默认情况下，每1000次迭代会在 data/snapshots 目录下保存一次快照。

本仓库实现了以下功能：

• 训练方式与 Caffe 实现相同——使用带有动量的 SGD，并结合 poly 学习率衰减策略。同时采用了权重衰减。最后一层的学习率是其他层的10倍。
• 实现了 Caffe 的 iter_size 参数，从而有效将批处理大小提升为 batch_size 乘以 iter_size。
• 在数据增强中使用了随机翻转和随机缩放。Caffe 实现使用4个固定尺度（0.5、0.75、1、1.25、1.5），而 PyTorch 实现则在每次迭代时随机选择[0.5,1.3]范围内的尺度。
• 当前版本中，边界标签（真值标签中的255）并未在损失函数中被忽略，而是与背景合并。未来版本中将实现 Caffe 的 ignore_label 参数。尚未实现使用 CRF 进行后处理。
• 训练过程中保持 BatchNorm 参数不变。同时，复现了 Caffe 中设置 use_global_stats = True 的行为。训练期间不计算运行均值和方差。

在 Nvidia Titan X GPU 上运行 train.py 时，大约会占用11.9 GB的显存。

评估

可以通过运行以下命令来评估已保存的模型：

python evalpyt.py

要查看每个命令行参数的说明，请运行：

python evalpyt.py -h

结果

在 PyTorch 中使用 MS COCO 预训练的初始化权重，并在 VOC 增强训练集（包含10582张图像）上进行训练后，验证集上的性能达到72.40%（evalpyt2.py，基于VOC数据）。相应的 .pth 文件可在此下载 链接。相比之下，使用 作者发布的 train_iter_20000.caffemodel 所获得的性能为75.54%，可通过运行 此文件 复现。同样，使用第一部分方法从 .caffemodel 转换而来的 .pth 模型也能达到75.54%的平均交并比。

此前的一个版本报告称，PyTorch 训练的模型在另一种计算方式下的平均交并比为78.48%（evalpyt.py，针对每张图像单独计算交并比，然后取平均值。这种计算平均交并比的方式与 作者所用的方法不同）。

要复现这一性能，请运行以下命令：

train.py --lr 0.00025 --wtDecay 0.0005 --maxIter 20000 --GTpath <此处填写训练真值图像路径> --IMpath <此处填写训练图像路径> --LISTpath data/list/train_aug.txt

数据集

结果部分展示的模型是在 这篇论文 发布的 VOC 增强训练集上训练的。您可以直接从 这里 下载该增强数据集。

请注意，这段代码也可用于训练适用于其他数据集的 pytorch-deeplab-resnet 模型。

致谢

部分代码借鉴自 https://github.com/ry/tensorflow-resnet。

pytorch-deeplab-resnet 快速上手指南

本指南基于 pytorch-deeplab-resnet 项目整理，旨在帮助开发者快速复现 DeepLab-ResNet 架构（源自 ACM MM 2017 论文 SketchParse）。该模型支持多尺度损失计算，适用于语义分割任务。

注意：本项目仅在 Python 2.7 环境下经过测试，请确保使用对应版本。

环境准备

系统要求

• 操作系统：Linux (推荐)
• Python 版本：2.7 (必须)
• GPU：NVIDIA GPU (推荐 Titan X 或更高，训练需约 11.9 GB 显存)
• CUDA/cuDNN：需安装与 PyTorch 版本匹配的驱动

前置依赖

1. PyTorch (Python 2.7 版本)
2. Caffe (DeepLab v2)：仅在进行权重转换时需要。
   • 仓库地址：deeplab-public-ver2
   • 需编译支持 Python 接口的 Caffe 版本。

安装步骤

1. 克隆代码库

git clone https://github.com/isht7/pytorch-deeplab-resnet.git
cd pytorch-deeplab-resnet

2. 准备初始化权重 (.pth)

模型需要使用在 MS COCO 数据集上预训练的权重进行初始化。你可以选择直接下载已转换的文件，或自行从 Caffe 模型转换。

方案 A：直接下载（推荐） 下载作者提供的已转换 .pth 文件：

• 初始化权重下载链接
• 下载后将其重命名为 init.caffemodel.pth (或根据代码实际读取的文件名调整) 并放入项目根目录。

方案 B：自行转换 如果你已有 init.caffemodel 文件，运行以下脚本进行转换（需确保 Caffe 和 PyTorch 环境已配置好）：

python init_net_surgery.py

基本使用

1. 训练模型 (Training)

使用 train.py 开始训练。默认配置会使用 SGD 优化器、动量策略以及 poly 学习率衰减策略。

最简单的训练命令：

python train.py --GTpath <训练集标签路径> --IMpath <训练集图片路径> --LISTpath data/list/train_aug.txt --NoLabels <类别数量>

参数说明：

• --GTpath: 地面真值（Ground Truth）图像的路径。
• --IMpath: 输入图像的路径。
• --LISTpath: 包含图像文件列表的 txt 文件路径（参考 data/list/ 下的格式）。
• --NoLabels: 训练数据中的类别总数。注意：标签值必须是连续的整数，从 0 开始（例如 7 个类别，标签应为 0,1,2,3,4,5,6）。边界标签（原为 255）在当前版本中会被合并到背景类。

高级训练示例（复现 78.48% mIOU 结果）：

python train.py --lr 0.00025 --wtDecay 0.0005 --maxIter 20000 --GTpath <train gt images path here> --IMpath <train images path here> --LISTpath data/list/train_aug.txt

训练过程中，模型快照每 1000 次迭代保存一次至 data/snapshots 目录。

2. 评估模型 (Evaluation)

训练完成后，使用 evalpyt.py 对保存的模型进行评估。

评估命令：

python evalpyt.py --modelPath <模型文件路径.pth> --GTpath <验证集标签路径> --IMpath <验证集图片路径> --LISTpath <验证集列表路径>

查看完整参数帮助：

python evalpyt.py -h

3. 数据集准备

官方结果是在 VOC augmented 训练集（10582 张图片）上得出的。

• 数据来源：VOC Augmented Data
• 该代码也支持使用自定义数据集进行训练，只需按照上述 --NoLabels 的要求整理标签即可。