PreciseRoIPooling

本仓库实现了在 ECCV 2018 大会上发表的论文《用于精确目标检测的定位置信度获取》中提出的 精确 RoI 池化（PrRoI Pooling）（口头报告）。

用于精确目标检测的定位置信度获取

江博睿*、罗瑞轩*、毛嘉源*、Tete Xiao、蒋宇宁 (* 表示共同第一作者。)

https://arxiv.org/abs/1807.11590

注意：为正确安装该库，请使用 git clone 克隆仓库，而非下载 ZIP 文件。这是因为 ./pytorch/prroi_pool/src/ 和 tensorflow/prroi_pool/src/kernels/external 文件夹内的源文件采用了符号链接的方式。直接下载 ZIP 文件会导致这些符号链接失效。此外，有报告指出 Windows 版本的 Git 也会破坏符号链接。详情请参阅 issues/58。

简介

简而言之，精确 RoI 池化是一种基于积分的平均池化方法，用于 RoI 池化操作。它避免了任何量化，并且对边界框坐标具有连续的梯度。具体来说：

• 它与原始的 RoI 池化（见 Fast R-CNN）不同。PrRoI Pooling 对每个 bin 使用平均池化而非最大池化，并且对边界框坐标具有连续的梯度。这意味着可以对某个损失函数关于每个 RoI 坐标的导数进行计算，从而优化 RoI 的位置。
• 它也不同于 Mask R-CNN 中提出的 RoI Align。PrRoI Pooling 使用完全基于积分的平均池化，而不是采样固定数量的点。这使得其关于坐标值的梯度是连续的。

为了更好地说明，我们在下图中对比展示了 RoI 池化、RoI Align 和 PrRoI Pooling。更多细节，包括梯度计算方法，可在我们的论文中找到。

实现

PrRoI Pooling 最初由 Tete Xiao 基于旷视科技内部开发的深度学习框架 MegBrain 实现。随后被移植到开源深度学习框架中。目前我们仅支持 PyTorch。遗憾的是，我们尚未制定将其实现移植到 TensorFlow 等其他框架的具体计划，但欢迎任何贡献（Pull Requests）。

使用方法（PyTorch 1.0）

在 pytorch/ 目录下，我们提供了基于 PyTorch 的 PrRoI Pooling 实现。它需要 PyTorch 1.0 及以上版本，并且仅支持 CUDA（未实现 CPU 模式）。 由于我们使用 PyTorch JIT 编译 C++/CUDA 代码，因此在代码中使用该模块时，只需执行以下操作：

from prroi_pool import PrRoIPool2D

avg_pool = PrRoIPool2D(window_height, window_width, spatial_scale)
roi_features = avg_pool(features, rois)

# 对于希望使用函数式接口的用户

from prroi_pool.functional import prroi_pool2d
roi_features = prroi_pool2d(features, rois, window_height, window_width, spatial_scale)

使用方法（PyTorch 0.4）

!!! 请先切换到 pytorch0.4 分支。

在 pytorch/ 目录下，我们提供了基于 PyTorch 的 PrRoI Pooling 实现。它需要 PyTorch 0.4，并且仅支持 CUDA（未实现 CPU 模式）。 要使用 PrRoI Pooling 模块，首先进入 pytorch/prroi_pool 目录并执行 ./travis.sh 来编译必要的组件（此步骤可能需要 nvcc）。然后在代码中使用该模块时，只需执行以下操作：

from prroi_pool import PrRoIPool2D

avg_pool = PrRoIPool2D(window_height, window_width, spatial_scale)
roi_features = avg_pool(features, rois)

# 对于希望使用函数式接口的用户

from prroi_pool.functional import prroi_pool2d
roi_features = prroi_pool2d(features, rois, window_height, window_width, spatial_scale)

其中：

• RoI 是一个形状为 m * 5 的浮点张量，格式为 (batch_index, x0, y0, x1, y1)，遵循原始 Caffe 实现中 RoI 池化的约定。不过，在某些框架中，批次索引是由整数张量提供的。
• spatial_scale 会乘以 RoI 的坐标。例如，如果您的特征图相对于输入图像进行了 16 倍下采样，则应使用 1/16 的空间尺度。
• RoI 的坐标遵循 [左闭右开] 的约定。即 (0, 0, 4, 4) 表示一个大小为 4x4 的矩形区域。

使用方法（TensorFlow）

在 tensorflow/ 目录下，我们提供了基于 TensorFlow 的 PrRoI Pooling 实现。经测试适用于 TensorFlow 2.2，且仅支持 CUDA（未实现 CPU 模式）。 要编译必要的组件，请按照以下说明操作：

要使用 PrRoI Pooling 模块，需先编译核心组件（此步骤可能需要 nvcc）。然后在代码中使用该模块时，只需执行以下操作：

需求

• CUDA 编译器 (NVCC)
• TensorFlow-GPU 2.x
• CMake
• Microsoft Visual C++ 构建工具（Windows 用户）

步骤说明

对于 Ubuntu 用户

CMake 配置

mkdir tensorflow/prroi_pool/build
cd tensorflow/prroi_pool/build
cmake -DCMAKE_BUILD_TYPE="Release" ..

编译并测试 PrRoI Pooling 模块

make

对于 Windows 用户

MSVC 配置

${MSVC_INSTALL_PATH}\VC\Auxiliary\Build\vcvars64.bat

CMake 配置

mkdir tensorflow/prroi_pool/build
cd tensorflow/prroi_pool/build
cmake -DCMAKE_BUILD_TYPE="Release" -G "NMake Makefiles" ..

编译并测试自定义算子

nmake BUILD=release

要在代码中使用该模块，只需执行：

from prroi_pool import PreciseRoIPooling

avg_pool = PreciseRoIPooling(window_height, window_width, spatial_scale, data_format)
roi_features = avg_pool([features, rois])

其中：

• RoI 是一个形状为 m * 5 的浮点张量，格式为 (batch_index, x0, y0, x1, y1)，遵循原始 Caffe 实现中 RoI 池化的约定。不过，在某些框架中，批次索引是由整数张量提供的。
• spatial_scale 会乘以 RoI 的坐标。例如，如果您的特征图相对于输入图像进行了 16 倍下采样，则应使用 1/16 的空间尺度。
• RoI 的坐标遵循 [左闭右开] 的约定。即 (0, 0, 4, 4) 表示一个大小为 4x4 的矩形区域。

PreciseRoIPooling 快速上手指南

PreciseRoIPooling (PrRoI Pooling) 是一种基于积分（双线性插值）的平均池化方法，旨在解决目标检测中定位置信度的问题。与传统的 RoI Pooling 和 RoI Align 不同，它避免了量化操作，并对边界框坐标具有连续梯度，支持直接对 RoI 坐标进行优化。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu) 或 Windows
• GPU: 必须拥有 NVIDIA GPU 并安装 CUDA 驱动（本项目不支持 CPU 模式）
• 编译器: nvcc (CUDA Compiler), gcc/g++ 或 MSVC (Windows)

前置依赖

根据你使用的深度学习框架选择对应环境：

PyTorch 用户:

• PyTorch 1.0+ (推荐) 或 PyTorch 0.4 (需切换分支)
• torchvision (可选，用于可视化)

TensorFlow 用户:

• TensorFlow-GPU 2.x (测试版本 2.2)
• CMake
• Windows 额外需要: Microsoft Visual C++ Build Tools

⚠️ 重要安装提示

请勿直接下载 ZIP 压缩包安装！ 由于源码中包含符号链接（symbolic links），直接下载 ZIP 会破坏链接结构导致编译失败。 请使用以下命令克隆仓库：

git clone https://github.com/vacancy/PreciseRoIPooling.git
cd PreciseRoIPooling

注：部分 Windows 版 Git 也可能破坏符号链接，如遇问题请参考官方 Issue #58。

---

安装步骤

方案 A：PyTorch 1.0+ (推荐)

PyTorch 1.0+ 版本利用 JIT 技术自动编译 C++/CUDA 代码，无需手动执行编译脚本，导入即可使用。

1. 确保已安装 PyTorch 1.0+ 且 CUDA 环境正常。
2. 进入 pytorch 目录（或在项目中直接引用该目录）。
3. 无需额外编译步骤，直接在代码中导入模块即可（见下文“基本使用”）。

方案 B：PyTorch 0.4

如果你必须使用 PyTorch 0.4，请先切换分支并手动编译。

1. 切换到对应分支：

   git checkout pytorch0.4
   

2. 进入目录并执行编译脚本（需要 nvcc）：

   cd pytorch/prroi_pool
   ./travis.sh
   

方案 C：TensorFlow 2.x

TensorFlow 版本需要手动编译自定义算子。

Ubuntu / Linux 用户

1. 创建构建目录并配置 CMake：

   mkdir tensorflow/prroi_pool/build
   cd tensorflow/prroi_pool/build
   cmake -DCMAKE_BUILD_TYPE="Release" ..
   

2. 编译并测试：

   make
   

Windows 用户

1. 配置 MSVC 环境变量（路径需根据实际情况调整）：

   ${MSVC_INSTALL_PATH}\VC\Auxiliary\Build\vcvars64.bat
   

2. 创建构建目录并配置 CMake（使用 NMake）：

   mkdir tensorflow\prroi_pool\build
   cd tensorflow\prroi_pool\build
   cmake -DCMAKE_BUILD_TYPE="Release" -G "NMake Makefiles" ..
   

3. 编译自定义算子：

   nmake BUILD=release
   

---

基本使用

PyTorch 使用示例

PrRoI Pooling 支持面向对象调用和函数式调用两种风格。

参数说明：

• rois: 形状为 (m, 5) 的浮点张量，格式为 (batch_index, x0, y0, x1, y1)。
• spatial_scale: 空间缩放比例。例如特征图相对于原图下采样了 16 倍，则该值为 1/16。
• 坐标遵循 [L, R) 惯例，即 (0, 0, 4, 4) 代表大小为 4x4 的框。

代码示例：

import torch
from prroi_pool import PrRoIPool2D

# 初始化参数
window_height = 7
window_width = 7
spatial_scale = 1.0 / 16.0

# 创建池化层
avg_pool = PrRoIPool2D(window_height, window_width, spatial_scale)

# 模拟输入数据
# features: (N, C, H, W)
features = torch.randn(2, 256, 56, 56).cuda()
# rois: (m, 5), format: (batch_idx, x0, y0, x1, y1)
rois = torch.tensor([
    [0, 10.5, 20.0, 50.0, 60.0],
    [1, 5.0, 10.0, 30.0, 40.0]
]).cuda()

# 前向传播
roi_features = avg_pool(features, rois)

# --- 或者使用函数式接口 ---
from prroi_pool.functional import prroi_pool2d
roi_features_func = prroi_pool2d(features, rois, window_height, window_width, spatial_scale)

TensorFlow 使用示例

代码示例：

import tensorflow as tf
from prroi_pool import PreciseRoIPooling

# 初始化参数
window_height = 7
window_width = 7
spatial_scale = 1.0 / 16.0
data_format = 'NHWC' # 或 'NCHW'

# 创建池化层
avg_pool = PreciseRoIPooling(window_height, window_width, spatial_scale, data_format)

# 模拟输入数据
features = tf.random.normal([2, 56, 56, 256])
rois = tf.constant([
    [0, 10.5, 20.0, 50.0, 60.0],
    [1, 5.0, 10.0, 30.0, 40.0]
], dtype=tf.float32)

# 前向传播 (输入为列表 [features, rois])
roi_features = avg_pool([features, rois])