一种更快速的 PyTorch 版 Faster R-CNN 实现

写在前面

[2020年5月29日] 本仓库大约在两年前启动，最初是首个支持多 GPU 训练的开源目标检测代码。它凝聚了众多开发者的心血。然而，近年来涌现出许多高质量的项目，例如：

• maskrcnn-benchmark
• detectron2
• mmdetection

目前，我认为该仓库在流程和代码风格上已经落后，因此将不再积极维护。尽管你仍然可以将其作为练习平台使用，但我强烈建议你转向上述项目，深入探索目标检测领域的前沿技术！

简介

:boom: 喜讯！本仓库现已支持 PyTorch 1.0！！！我们借鉴了 maskrcnn-benchmark 中的部分代码和技巧。请切换到 pytorch-1.0 分支！

本项目是一个 更快 的 PyTorch 版 Faster R-CNN 实现，旨在加速 Faster R-CNN 目标检测模型的训练。近期已有不少优秀的实现：

• rbgirshick/py-faster-rcnn，基于 PyCaffe + NumPy 开发
• longcw/faster_rcnn_pytorch，基于 PyTorch + NumPy 开发
• endernewton/tf-faster-rcnn，基于 TensorFlow + NumPy 开发
• ruotianluo/pytorch-faster-rcnn，基于 PyTorch + TensorFlow + NumPy 开发

我们在实现过程中参考了以上项目，尤其是 longcw/faster_rcnn_pytorch。不过，我们的实现相比这些项目有几项独特且创新的功能：

• 纯 PyTorch 代码。我们将所有 NumPy 实现转换为 PyTorch。
• 支持多图像批量训练。我们对数据加载器、RPN、RoI 池化等各个模块进行了改造，使其能够在一个 mini-batch 中处理多张图像。
• 支持多 GPU 训练。通过使用 nn.DataParallel 多 GPU 封装，结合上述两项特性，我们可以灵活地选择使用单 GPU 或多 GPU 进行训练。
• 支持三种池化方法。我们集成了 RoI Pooling、RoI Align 和 RoI Crop 三种池化方式，并对其进行了修改以支持多图像批量训练。
• 内存高效。我们对图像的宽高比进行限制，将相似宽高比的图像分组到同一个 mini-batch 中。这样，在单块 Titan X（12 GB 显存）上，我们就可以以 batch size = 4 的规模训练 ResNet101 和 VGG16；而在 8 卡环境下，每张卡的最大 batch size 可达 3（ResNet101），总 batch size 为 24。
• 速度更快。基于上述优化，训练速度显著提升。我们在下表中展示了在 NVIDIA TITAN Xp 上的训练速度。

我们正在做以及未来计划做的事情

• 同时支持 Python 2 和 Python 3（特别感谢 cclauss）。
• 添加可变形池化层（主要由 Xander 支持）。
• 支持 PyTorch 0.4.0（此分支）。
• 支持 tensorboardX。
• 支持 PyTorch 1.0（请切换到 pytorch-1.0 分支）。

其他相关实现

• 特征金字塔网络 (FPN)

• Mask R-CNN（正在进行中 已由 roytseng-tw 实现）

• 图 R-CNN（用于场景图生成的扩展）

教程

• 博客 由 ankur6ue 撰写

基准测试

我们在三个数据集上对我们的代码进行了全面的基准测试：Pascal VOC、COCO 和 Visual Genome，使用了两种不同的网络架构：VGG16 和 ResNet101。以下是结果：

1). PASCAL VOC 2007（训练/测试：07trainval/07test，尺度=600，ROI Align）

模型	GPU 数量	批量大小	学习率	学习率衰减	最大 epoch 数	每 epoch 时间	每 GPU 显存	mAP
VGG-16	1	1	1e-3	5	6	0.76 小时	3265MB	70.1
VGG-16	1	4	4e-3	8	9	0.50 小时	9083MB	69.6
VGG-16	8	16	1e-2	8	10	0.19 小时	5291MB	69.4
VGG-16	8	24	1e-2	10	11	0.16 小时	11303MB	69.2
Res-101	1	1	1e-3	5	7	0.88 小时	3200 MB	75.2
Res-101	1	4	4e-3	8	10	0.60 小时	9700 MB	74.9
Res-101	8	16	1e-2	8	10	0.23 小时	8400 MB	75.2
Res-101	8	24	1e-2	10	12	0.17 小时	10327MB	75.1

2). COCO（训练/测试：coco_train+coco_val-minival/minival，尺度=800，最大尺寸=1200，ROI Align）

模型	GPU 数量	批量大小	学习率	学习率衰减	最大 epoch 数	每 epoch 时间	每 GPU 显存	mAP
VGG-16	8	16	1e-2	4	6	4.9 小时	7192 MB	29.2
Res-101	8	16	1e-2	4	6	6.0 小时	10956 MB	36.2
Res-101	8	16	1e-2	4	10	6.0 小时	10956 MB	37.0

注意。由于上述模型使用尺度=800，您需要在测试命令末尾添加“--ls”。

3). COCO（训练/测试：coco_train+coco_val-minival/minival，尺度=600，最大尺寸=1000，ROI Align）

模型	GPU 数量	批量大小	学习率	学习率衰减	最大 epoch 数	每 epoch 时间	每 GPU 显存	mAP
Res-101	8	24	1e-2	4	6	5.4 小时	10659 MB	33.9
Res-101	8	24	1e-2	4	10	5.4 小时	10659 MB	34.5

4). Visual Genome（训练/测试：vg_train/vg_test，尺度=600，最大尺寸=1000，ROI Align，类别=2500）

模型	GPU 数量	批量大小	学习率	学习率衰减	最大 epoch 数	每 epoch 时间	每 GPU 显存	mAP
VGG-16	1 P100	4	1e-3	5	20	3.7 小时	12707 MB	4.4

感谢 Remi 提供的 Visual Genome 上的预训练检测模型！

• 点击上述表格中的链接以下载我们预训练的 Faster R-CNN 模型。
• 如果未特别说明，我们使用的 GPU 是 NVIDIA Titan X Pascal（12GB）。

准备工作

首先，克隆代码

git clone https://github.com/jwyang/faster-rcnn.pytorch.git

然后，创建一个文件夹：

cd faster-rcnn.pytorch && mkdir data

前提条件

• Python 2.7 或 3.6
• PyTorch 0.4.0（目前不支持 0.4.1 或更高版本）
• CUDA 8.0 或更高版本

数据准备

• PASCAL_VOC 07+12：请按照 py-faster-rcnn 中的说明准备 VOC 数据集。实际上，您可以参考其他任何方法。下载数据后，在 data/ 文件夹中创建软链接。

• COCO：同样，请按照 py-faster-rcnn 中的说明准备数据。

• Visual Genome：请按照 bottom-up-attention 中的说明准备 Visual Genome 数据集。您需要先下载图像和目标标注文件，然后根据该仓库提供的脚本进行预处理，以获得词汇表和清理后的标注。

预训练模型

我们在实验中使用了两个预训练模型，VGG 和 ResNet101。您可以从以下链接下载这两个模型：

• VGG16：Dropbox，VT Server

• ResNet101：Dropbox，VT Server

下载后将其放入 data/pretrained_model/ 目录中。

注意。我们比较了 PyTorch 和 Caffe 的预训练模型，令人惊讶的是，Caffe 的预训练模型性能略优于 PyTorch 的预训练模型。我们建议使用上述链接中的 Caffe 预训练模型来复现我们的结果。

如果您想使用 PyTorch 预训练模型，请务必先将图像从 BGR 转换为 RGB，并且使用与预训练模型相同的预处理方式（减去均值并归一化）。

编译

如 ruotianluo/pytorch-faster-rcnn 所述，在 make.sh 文件中选择正确的 -arch 参数，以编译 CUDA 代码：

GPU 型号	架构
TitanX (Maxwell/Pascal)	sm_52
GTX 960M	sm_50
GTX 1080 (Ti)	sm_61
Grid K520 (AWS g2.2xlarge)	sm_30
Tesla K80 (AWS p2.xlarge)	sm_37

有关设置架构的更多详细信息，请参阅 此处 或 此处。

使用 pip 安装所有 Python 依赖项：

pip install -r requirements.txt

通过以下简单命令编译 CUDA 依赖项：

cd lib
sh make.sh

这将编译您所需的所有模块，包括 NMS、ROI_Pooling、ROI_Align 和 ROI_Crop。默认版本是使用 Python 2.7 编译的，如果您使用的是其他 Python 版本，请自行编译。

正如此 issue 中所指出的那样，如果在编译过程中遇到错误，可能是您未将 CUDA 路径导出到环境变量中。

训练

在训练之前，请设置正确的目录来保存和加载训练好的模型。请修改 trainval_net.py 和 test_net.py 中的 save_dir 和 load_dir 参数，以适应您的环境。

要在 Pascal VOC 数据集上使用 VGG16 训练 Faster R-CNN 模型，只需运行：

CUDA_VISIBLE_DEVICES=$GPU_ID python trainval_net.py \
                   --dataset pascal_voc --net vgg16 \
                   --bs $BATCH_SIZE --nw $WORKER_NUMBER \
                   --lr $LEARNING_RATE --lr_decay_step $DECAY_STEP \
                   --cuda

其中 bs 是批大小，默认为 1。或者，要在 Pascal VOC 数据集上使用 ResNet101 训练，只需运行：

CUDA_VISIBLE_DEVICES=$GPU_ID python trainval_net.py \
                    --dataset pascal_voc --net res101 \
                    --bs $BATCH_SIZE --nw $WORKER_NUMBER \
                    --lr $LEARNING_RATE --lr_decay_step $DECAY_STEP \
                    --cuda

上述 BATCH_SIZE 和 WORKER_NUMBER 可根据您的 GPU 显存大小进行调整。在显存为 12G 的 Titan Xp 上，最大可设为 4。

如果您有多个（比如 8 个）Titan Xp GPU，那就全部用起来吧！尝试以下命令：

python trainval_net.py --dataset pascal_voc --net vgg16 \
                       --bs 24 --nw 8 \
                       --lr $LEARNING_RATE --lr_decay_step $DECAY_STEP \
                       --cuda --mGPUs

如果想在 COCO 或 Visual Genome 数据集上训练，可将 dataset 参数改为 “coco” 或 “vg”。

测试

如果您想评估一个预训练的 VGG16 模型在 Pascal VOC 测试集上的检测性能，只需运行：

python test_net.py --dataset pascal_voc --net vgg16 \
                   --checksession $SESSION --checkepoch $EPOCH --checkpoint $CHECKPOINT \
                   --cuda

请指定具体的模型会话、检查周期和检查点，例如 SESSION=1，EPOCH=6，CHECKPOINT=416。

演示

如果您想使用预训练模型对自己的图像进行目标检测，可以先下载上述表格中列出的预训练模型，或者先训练自己的模型，然后将图像放入 $ROOT/images 文件夹中，再运行：

python demo.py --net vgg16 \
               --checksession $SESSION --checkepoch $EPOCH --checkpoint $CHECKPOINT \
               --cuda --load_dir path/to/model/directoy

检测结果将出现在 $ROOT/images 文件夹中。

请注意，默认的 demo.py 仅支持 Pascal VOC 类别。您需要修改 此行，以适配您自己的模型。

以下是一些检测结果：

网络摄像头演示

您可以通过运行以下命令，在实时演示中使用网络摄像头：

python demo.py --net vgg16 \
               --checksession $SESSION --checkepoch $EPOCH --checkpoint $CHECKPOINT \
               --cuda --load_dir path/to/model/directoy \
               --webcam $WEBCAM_ID

演示可通过点击图像窗口并按下 ‘q’ 键来停止。

作者

该项目由 Jianwei Yang 和 Jiasen Lu 以及其他许多人共同贡献（感谢他们！）。

引用

@article{jjfaster2rcnn,
    Author = {Jianwei Yang and Jiasen Lu and Dhruv Batra and Devi Parikh},
    Title = {一种更快的 Faster R-CNN PyTorch 实现},
    Journal = {https://github.com/jwyang/faster-rcnn.pytorch},
    Year = {2017}
}

@inproceedings{renNIPS15fasterrcnn,
    Author = {Shaoqing Ren and Kaiming He and Ross Girshick and Jian Sun},
    Title = {Faster {R-CNN}: 向基于区域建议网络的实时目标检测迈进},
    Booktitle = {神经信息处理系统进展 ({NIPS})},
    Year = {2015}
}

faster-rcnn.pytorch 快速上手指南

⚠️ 重要提示：本项目已停止主动维护，代码风格和流程相对过时。仅建议作为学习参考或实验场（Playground）。对于生产环境或深入研究，强烈推荐使用 Detectron2 或 MMDetection 等现代框架。本指南基于原文提供的 pytorch-1.0 分支及经典用法整理。

1. 环境准备

在开始之前，请确保满足以下系统和依赖要求：

• 操作系统: Linux (推荐)
• Python: 2.7 或 3.6
• PyTorch:
  • 主分支支持 PyTorch 0.4.0 (不支持 0.4.1 及以上版本)。
  • 若需使用新版，请切换至 pytorch-1.0 分支。
• CUDA: 8.0 或更高版本
• GPU: 支持 CUDA 的 NVIDIA 显卡（编译时需指定对应的架构代码）

2. 安装步骤

2.1 克隆代码与创建目录

git clone https://github.com/jwyang/faster-rcnn.pytorch.git
cd faster-rcnn.pytorch && mkdir data

2.2 安装 Python 依赖

建议使用国内镜像源加速安装：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

2.3 编译 CUDA 扩展

根据显卡型号修改 lib/make.sh 中的 -arch 参数（例如 Titan X 为 sm_52, GTX 1080 为 sm_61），然后执行编译：

cd lib
sh make.sh

注：此步骤将编译 NMS, ROI_Pooling, ROI_Align 和 ROI_Crop 等模块。如遇报错，请检查是否已正确导出 CUDA 环境变量。

2.4 准备预训练模型

下载骨干网络预训练模型并放入 data/pretrained_model/ 目录。原文推荐使用 Caffe 转制的模型以获得更好效果：

• VGG16: 下载链接
• ResNet101: 下载链接

# 示例：将下载的 .pth 文件移动至指定目录
mv vgg16_caffe.pth data/pretrained_model/

2.5 数据集准备

按照 py-faster-rcnn 的说明准备 PASCAL VOC 或 COCO 数据集，并在 data/ 目录下建立软链接。

3. 基本使用

训练模型

以在 PASCAL VOC 数据集上使用 VGG16 backbone 进行训练为例：

CUDA_VISIBLE_DEVICES=$GPU_ID python trainval_net.py \
                   --dataset pascal_voc --net vgg16 \
                   --bs $BATCH_SIZE --nw $WORKER_NUMBER \
                   --lr $LEARNING_RATE --lr_decay_step $DECAY_STEP \
                   --cuda

参数说明：

• $GPU_ID: 指定使用的 GPU 编号（如 0）。
• --net: 选择网络架构，可选 vgg16 或 res101。
• --bs: Batch Size，根据显存大小调整（单卡 Titan X 可设 4）。
• --cuda: 启用 GPU 训练。

若使用 ResNet101，只需将 --net 参数改为 res101：

CUDA_VISIBLE_DEVICES=$GPU_ID python trainval_net.py \
                    --dataset pascal_voc --net res101 \
                    --bs $BATCH_SIZE --nw $WORKER_NUMBER \
                    --lr $LEARNING_RATE --lr_decay_step $DECAY_STEP \
                    --cuda

测试模型

训练完成后，使用 test_net.py 进行评估（具体命令参考原库文档，需注意不同数据集的缩放比例参数，如 COCO 数据集测试时可能需要添加 --ls 参数）。