batched_nms() 和 nms()内置区别在于

torchvision.ops.nms

不区分类别对所有bbox进行过滤。如果有不同类别的bbox重叠的话会导致被过滤掉并不会分开计算。  

boxes (Tensor[N, 4])) – bounding boxes坐标. 格式：(x1, y1, x2, y2)
scores (Tensor[N]) – bounding boxes得分
iou_threshold (float) – IoU过滤阈值
返回NMS过滤后的bouding boxes索引（降序排列）
 

torchvision.ops.batched_nms

根据每个类别进行过滤，只对同一种类别进行计算IOU和阈值过滤。

boxes: Tensor, 预测框
scores: Tensor, 预测置信度
idxs: Tensor, 预测框类别
iou_threshold: float, IOU阈值

demo

import torchvision.ops as ops
import torch
 
 
b = torch.Tensor([[2,2,4,4], [1,1,5,5], [3,3,3.5,3.9]]) # bbox
c = torch.Tensor([0,1,0]) # classes
s = torch.Tensor([0.8,0.8,0.8]) # scores
 
 
 
ops.batched_nms(b, s, c, 0.001)
#运行结果 tensor([1, 2])
#[1,1,5,5], [3,3,3.5,3.9] bbox实际上是有包含关系的，但是类别不一样
 
 
ops.nms(b, s, 0.001)
# 运行结果 tensor([0])
# 可以看到 [1,1,5,5] 类别为1 但是被过滤掉了，只留下0号类别的[2,2,4,4]