（beta）torch_npu.contrib.function.npu_batched_multiclass_nms

接口原型

torch_npu.contrib.function.npu_batched_multiclass_nms(multi_bboxes,multi_scores,score_thr=0.05,nms_thr=0.45,max_num=50,score_factors=None):

功能描述

使用NPU API的批量多类bbox NMS。

参数说明

• multi_bboxes (Tensor) - shape(bs, n, #class, 4)或(bs, n, 4)。
• multi_scores (Tensor) - shape(bs, n, #class+1)，其中最后一列包含background class分数，可忽略。在NPU上，为保持语义畅通，我们将统一维度。
• score_thr (Float，默认值为0.05) - bbox阈值，分数低于它的bbox将不被考虑。
• nms_thr (Float，默认值为0.45) - NMS IoU阈值。最初的实现是传递{"iouthreshold": 0.45}字典，这里做了简化。
• max_num (Int，默认值为50) - 如果NMS后的bbox数超过max_num值，则只保留最大max_num；如果NMS后的bbox数小于max_num值，则输出将零填充到max_num值。在NPU上需提前申请内存，因此目前不能将max_num值设置为-1。
• score_factors (Tensor，默认值为None) - NMS应用前用来乘分数的因子。

输出说明

Tuple - (bboxes, labels)，shape为(bs, k, 5)和(bs, k, 1)的张量。标签以0为基础。

约束说明

在动态shape条件下，由于NPU op的限制，最多支持20个类别（nmsed_classes）和10000个框（nmsed_boxes）。

支持的型号

• Atlas 训练系列产品
• Atlas A2 训练系列产品
• Atlas A3 训练系列产品
• Atlas 推理系列产品

调用示例

1
2
3
4

>>> from torch_npu.contrib.function import npu_batched_multiclass_nms
>>> boxes = torch.randint(1, 255, size=(4, 200, 80, 4)).npu().half()
>>> scores = torch.randn(4, 200, 81).npu().half()
>>> det_bboxes, det_labels = npu_batched_multiclass_nms(boxes, scores, score_thr=0.3, nms_thr=0.5, max_num=3)