RoIAlign for PyTorch

这是一个 RoIAlign 的 PyTorch 实现。 该实现基于 crop_and_resize，并在 CPU 和 GPU 上同时支持前向和反向传播。

注意： 感谢 meikuam 为 PyTorch 1.0 更新了此仓库。您可以在 pytorch_0.4 分支中找到适用于 torch <= 0.4.1 的原始版本。

简介

crop_and_resize 函数源自 tensorflow，其接口与 TensorFlow 版本相同，唯一的区别是 PyTorch 中的输入特征图应采用 NCHW 格式。 正如我们所预期的那样，它们在前向和反向传播中的输出值也完全一致（误差 < 1e-5），具体比较请参见 test.py。

注： crop_and_resize 的文档可以在这里找到 这里。 而 RoIAlign 是对 crop_and_resize 的封装，它接受未归一化的 (x1, y1, x2, y2) 坐标框作为输入（而 crop_and_resize 则使用归一化的 (y1, x1, y2, x2) 坐标）。 有关 RoIAlign 和 crop_and_resize 之间差异的更多详细信息，请参阅 tensorpack。

警告： 目前仅支持使用默认的 GPU（索引 0）。

使用方法

• 安装并测试

  python setup.py install
  ./test.sh
  

• 使用 RoIAlign 或 crop_and_resize

  自 PyTorch 1.2.0 起，具有非静态前向方法的传统自动求导函数已被弃用。 我们现在使用具有静态前向方法的新式自动求导函数。示例如下：

  import torch
  from roi_align import RoIAlign      # RoIAlign 模块
  from roi_align import CropAndResize # crop_and_resize 模块
  
  # 输入特征图（假设 batch_size==2）
  image = torch.arange(0., 49).view(1, 1, 7, 7).repeat(2, 1, 1, 1)
  image[0] += 10
  print('image: ', image)
  
  
  # 例如，我们有两个 bbox，坐标格式为 xyxy（第一个属于 batch_id=0，第二个属于 batch_id=1）。
  boxes = torch.Tensor([[1, 0, 5, 4],
                       [0.5, 3.5, 4, 7]])
  
  box_index = torch.tensor([0, 1], dtype=torch.int) # 每个 bbox 在 batch 中的索引
  
  # RoIAlign 层，设置裁剪大小：
  crop_height = 4
  crop_width = 4
  roi_align = RoIAlign(crop_height, crop_width)
  
  # 进行裁剪：
  crops = roi_align(image, boxes, box_index)
  
  print('crops:', crops)
  

  输出：

  image:  tensor([[[[10., 11., 12., 13., 14., 15., 16.],
        [17., 18., 19., 20., 21., 22., 23.],
        [24., 25., 26., 27., 28., 29., 30.],
        [31., 32., 33., 34., 35., 36., 37.],
        [38., 39., 40., 41., 42., 43., 44.],
        [45., 46., 47., 48., 49., 50., 51.],
        [52., 53., 54., 55., 56., 57., 58.]]],
  
  
      [[[ 0.,  1.,  2.,  3.,  4.,  5.,  6.],
        [ 7.,  8.,  9., 10., 11., 12., 13.],
        [14., 15., 16., 17., 18., 19., 20.],
        [21., 22., 23., 24., 25., 26., 27.],
        [28., 29., 30., 31., 32., 33., 34.],
        [35., 36., 37., 38., 39., 40., 41.],
        [42., 43., 44., 45., 46., 47., 48.]]]])
        
  crops: tensor([[[[11.0000, 12.0000, 13.0000, 14.0000],
            [18.0000, 19.0000, 20.0000, 21.0000],
            [25.0000, 26.0000, 27.0000, 28.0000],
            [32.0000, 33.0000, 34.0000, 35.0000]]],
  
  
          [[[24.5000, 25.3750, 26.2500, 27.1250],
            [30.6250, 31.5000, 32.3750, 33.2500],
            [36.7500, 37.6250, 38.5000, 39.3750],
            [ 0.0000,  0.0000,  0.0000,  0.0000]]]])
  

RoIAlign.pytorch 快速上手指南

环境准备

• 系统要求：Linux / macOS（Windows 需自行配置编译环境）
• Python 版本：建议 Python 3.6+
• PyTorch 版本：支持 PyTorch 1.0+（本指南基于新版静态 forward 方法）
• 硬件要求：
  • 支持 CPU 和 GPU 推理与反向传播
  • 注意：当前版本仅默认支持单卡运行（GPU index 0）
• 依赖项：
  • torch
  • setup.py 构建工具（通常随 PyTorch 安装自带）

💡 国内用户建议使用清华或阿里镜像源加速 PyTorch 安装：

pip install torch -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

克隆仓库并安装模块：

git clone https://github.com/longcw/RoIAlign.pytorch.git
cd RoIAlign.pytorch
python setup.py install

验证安装是否成功：

./test.sh

若脚本无报错且输出对比结果误差小于 1e-5，则安装成功。

基本使用

以下是最小可用示例，演示如何使用 RoIAlign 进行特征图裁剪与对齐：

import torch
from roi_align import RoIAlign

# 构造输入特征图 (batch_size=2, channel=1, H=7, W=7)
image = torch.arange(0., 49).view(1, 1, 7, 7).repeat(2, 1, 1, 1)
image[0] += 10  # 区分两个 batch 的数据

# 定义边界框 (x1, y1, x2, y2)，未归一化坐标
boxes = torch.Tensor([[1, 0, 5, 4],      # 属于 batch 0
                      [0.5, 3.5, 4, 7]]) # 属于 batch 1

# 指定每个 box 所属的 batch 索引
box_index = torch.tensor([0, 1], dtype=torch.int)

# 初始化 RoIAlign 层，设定输出裁剪尺寸
crop_height = 4
crop_width = 4
roi_align = RoIAlign(crop_height, crop_width)

# 执行前向传播
crops = roi_align(image, boxes, box_index)

print('crops:', crops)

该代码将输出两个经过双线性插值对齐后的 4x4 特征图块，可直接用于后续检测头网络。