torch_npu.npu_nms_with_mask

API接口

torch_npu.npu_nms_with_mask(input, iou_threshold) -> (Tensor, Tensor, Tensor)

功能描述

生成值0或1，用于nms算子确定有效位。

参数说明

• input (Tensor) - 输入张量
• iou_threshold (Scalar) - 阈值。如果超过此阈值，则值为1，否则值为0。

输出说明

• selected_boxes (Tensor) - shape为[N,5]的2D张量，表示filtered box，包括proposal box和相应的置信度分数。
• selected_idx (Tensor) - shape为[N]的1D张量，表示输入建议框的index。
• selected_mask (Tensor) - shape为[N]的1D张量，判断输出建议框是否有效。

约束说明

输入box_scores的2nd-dim必须等于8。

示例

>>> input = torch.tensor([[0.0, 1.0, 2.0, 3.0, 0.6], [6.0, 7.0, 8.0, 9.0, 0.4]], dtype=torch.float16).to("npu")
>>> iou_threshold = 0.5
>>> output1, output2, output3, = torch_npu.npu_nms_with_mask(input, iou_threshold)
>>> output1
tensor([[0.0000, 1.0000, 2.0000, 3.0000, 0.6001],
        [6.0000, 7.0000, 8.0000, 9.0000, 0.3999]], device='npu:0',      dtype=torch.float16)
>>> output2
tensor([0, 1], device='npu:0', dtype=torch.int32)
>>> output3
tensor([1, 1], device='npu:0', dtype=torch.uint8)