TopK

从输入张量中检索前 k 个元素，并输出包含这些值的张量和包含其索引的另一个张量。

属性

operation TopK 操作可以是以下之一

• MIN 检索前最小元素。

• MAX 检索前最大元素。

axes 采样的轴；必须是最后四个维度之一。

k 要检索的元素数量。对于维度大小为 \(d\) 的 axis，k 必须遵守 \(k<=d, k<=3840\)。

输入

input: T1 类型的张量。

输出

values: T1 类型的张量。

indices: T2 类型的张量。

数据类型

T1: int32, int64, float16, float32, bfloat16

T2: int32

形状信息

input 是形状为 \([a_0,...,a_n]\) 的张量

values 是形状为 \([a_0,...,a_{axes-1},k,a_{axes+1},...,a_n]\) 的张量

indices 是形状为 \([a_0,...,a_{axes-1},k,a_{axes+1},...,a_n]\) 的张量

体积限制

input、output 和 indices 最多可以有 \(2^{31}-1\) 个元素。

示例

TopK

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 5))
layer = network.add_topk(in1, op=trt.TopKOperation.MAX, k=3, axes=2)
network.mark_output(layer.get_output(0))
network.mark_output(layer.get_output(1))

inputs[in1.name] = np.array([[-3.0, -2.0, -1.0, 10.0, -25.0], [0.0, 1.0, 2.0, -2.0, -1.0]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[10.0, -1.0, -2.0], [2.0, 1.0, 0.0]])

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

C++ API

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

Python API

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