TopK#
Retrieves the top k
elements from an input tensor and outputs a tensor with those values and another tensor with their indices.
Attributes#
operation
TopK operation can be one of:
MIN
Retrieves the top minimum elements.MAX
Retrieves the top maximum elements.
axes
the axis to sample; must be one of the last four dimensions.
k
The number of elements to retrieve. For axis
with dimension size \(d\), k
must obey \(k<=d, k<=3840\).
Inputs#
input: tensor of type T1
.
Shape Information#
input is a tensor with a shape of \([a_0,...,a_n]\)
values is a tensor of shape \([a_0,...,a_{axes-1},k,a_{axes+1},...,a_n]\)
indices is a tensor of shape \([a_0,...,a_{axes-1},k,a_{axes+1},...,a_n]\)
Volume Limits#
input, output, and indices can have up to \(2^{31}-1\) elements.
Examples#
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#
For more information about the C++ ITopKLayer operator, refer to the C++ ITopKLayer documentation.
Python API#
For more information about the Python ITopKLayer operator, refer to the Python ITopKLayer documentation.