PyTorch学习之OPS (NMS, RoIAlign, RoIPool)

位于torchvision.ops下（pytorch>=1.2.0， torchvision >= 0.3）

NMS:

torchvision.ops.nms(boxes, scores, iou_threshold)

参数：

• boxes (Tensor[N, 4])) – bounding boxes坐标. 格式：(x1, y1, x2, y2)
• scores (Tensor[N]) – bounding boxes得分
• iou_threshold (float) – IoU过滤阈值

返回值：

• keep :NMS过滤后的bouding boxes索引（降序排列）

RoIAlign: 用于Mask R-CNN

torchvision.ops.roi_align(input, boxes, output_size, spatial_scale=1.0, sampling_ratio=-1)

参数：

• input (Tensor[N, C, H, W]) – 输入张量
• boxes (Tensor[K, 5] or List[Tensor[L, 4]]) – 输入的box 坐标，格式：list(x1, y1, x2, y2) 或者(batch_index, x1, y1, x2, y2)
• output_size (int or Tuple[int, int]) – 输出尺寸, 格式： (height, width)
• spatial_scale (float) – 将输入坐标映射到box坐标的尺度因子. 默认: 1.0
• sampling_ratio (int) – 插值网格上用来计算池化后输出的采样点数量；如果sampling_ratio>0, sampling_ratio个采样点将会被使用，如果sampling_ratio<= 0，自适应采样点数量,即ceil(roi_width / pooled_w)和ceil(roi_height / pooled_h)，默认： sampling_ratio =-1

RoIPool:用于Fast R-CNN

torchvision.ops.roi_pool(input, boxes, output_size, spatial_scale=1.0)

参数：

• input (Tensor[N, C, H, W]) – 输入张量
• boxes (Tensor[K, 5] or List[Tensor[L, 4]]) – 输入的box 坐标，格式：list(x1, y1, x2, y2) 或者(batch_index, x1, y1, x2, y2)
• output_size (int or Tuple[int, int]) – 输出尺寸, 格式： (height, width)
• spatial_scale (float) – 将输入坐标映射到box坐标的尺度因子. 默认: 1.0