NMS¶

NMS 算法遍历一组边界框及其置信度分数，分数按降序排列。如果框的分数高于给定的阈值，并且它们与先前选择的框的交并比 (IoU) 小于或等于给定的阈值，则会选择这些框。此层实现每个批次项目和每个类别的 NMS。

对于每个批次项目，框最初按其分数排序，而不考虑类别。仅考虑最多 TopK 限制的框进行选择（每个批次）。在选择期间，仅比较同一类别的重叠框，以便不同类别的重叠框不会相互抑制。

对于每个批次项目，具有相同分数的候选边界框的顺序是未指定的，但对于相同的输入，不同运行之间的顺序将是一致的。

属性¶

fmt 边界框格式可以是以下之一

CORNER_PAIRS (x1, y1, x2, y2)，其中 (x1, y1) 和 (x2, y2) 是任意一对对角角。
CENTER_SIZES (x_center, y_center, width, height)，其中 (x_center, y_center) 是框的中心点。

默认值为 CORNER_PAIRS。

limit TopK 框限制是每个批次项目考虑选择的最大过滤框数。SM 5.3 和 6.2 设备的默认值为 2000，否则为 5000。TopK 框限制必须小于或等于 {SM 5.3 和 6.2 设备的 2000，否则为 5000}。

输入¶

Boxes: 类型为 T 的张量。

Scores: 类型为 T 的张量。

MaxOutputBoxesPerClass: 类型为 int32 的张量。

IoUThreshold 可选: 类型为 float32 的张量。这必须包含范围 [0.0f, 1.0f] 中的标量值。阈值的默认值为 0.0f。

ScoreThreshold 可选: 类型为 float32 的张量。阈值的默认值为 0.0f。

输出¶

SelectedIndices: 类型为 int32 的张量。

NumOutputBoxes: 类型为 int32 的张量。

数据类型¶

T: float16, float32, bfloat16

形状信息¶

Boxes 是一个张量，如果框是按类别划分的，则形状为 [batchSize, numInputBoundingBoxes, numClasses, 4]，如果相同的框用于每个类别，则形状为 [batchSize, numInputBoundingBoxes, 4]。

Scores 是一个张量，形状为 [batchSize, numInputBoundingBoxes, numClasses]。

MaxOutputBoxesPerClass 是一个包含标量的 0D 张量。

IoUThreshold 是一个包含标量的 0D 张量。

ScoreThreshold 是一个包含标量的 0D 张量。

SelectedIndices 是一个张量，形状为 [NumOutputBoxes, 3]。每行包含一个元组 (batchIndex, classIndex, boxIndex)。

NumOutputBoxes 是一个包含标量的 0D 张量。

容量限制¶

Boxes、Scores 和 SelectedIndices 最多可以有 \(2^{31}-1\) 个元素。

示例¶

非极大值抑制 (NMS)

opt_profile = get_runner.builder.create_optimization_profile()
get_runner.config.add_optimization_profile(opt_profile)

boxes = network.add_input("boxes", dtype=trt.float32, shape=(1, 3, 4))
scores = network.add_input("scores", dtype=trt.float32, shape=(1, 3, 3))
constant = network.add_constant(shape=(), weights=np.ones(shape=(), dtype=np.int32))
max_output_boxes_per_class = constant.get_output(0)

layer = network.add_nms(boxes, scores, max_output_boxes_per_class)
network.mark_output(layer.get_output(0))
network.mark_output(layer.get_output(1))
layer.get_output(0).dtype = trt.int32
layer.get_output(1).dtype = trt.int32

inputs[boxes.name] = np.array([[[0.0, 0.0, 0.1, 0.1], [0.2, 0.2, 0.4, 0.4], [0.5, 0.5, 0.6, 0.6]]])
inputs[scores.name] = np.array([[[0.7, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.9]]])

# Expected shape is [2, 3]
outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[0, 2, 2], [0, 0, 0]])

# Expected shape is [] with a scalar value of 2
outputs[layer.get_output(1).name] = layer.get_output(1).shape
expected[layer.get_output(1).name] = np.array(2)

C++ API¶

有关 C++ INMSLayer 算子的更多信息，请参阅 C++ INMSLayer 文档。

Python API¶

有关 Python INMSLayer 算子的更多信息，请参阅 Python INMSLayer 文档。