Gather¶
将输入张量的元素收集到输出张量中。
属性¶
axis 从中收集元素的轴，必须遵守 \(0 \leq axis < rank(input)\)。
mode Gather 模式
DEFAULT类似于 ONNX Gather。这是默认模式。ELEMENT类似于 ONNX GatherElements。ND类似于 ONNX GatherND。
num_elementwise_dims 开始收集的维度。
输入¶
input: T1 类型的张量。
indices: T2 类型的张量。 indices 中的每个元素 \(I_j\) 必须遵守 \(0 \leq I_j < \text{input_dimensions}[axis]\)
输出¶
output: T1 类型的张量。
数据类型¶
T1: bool, int4, int8, int32, int64, float8, float16, float32, bfloat16
T2: int32, int64
如果 indices 是
int64且数据不是int64，则 indices 将被转换为int32。
形状信息¶
input 是一个形状为 \([a_0,...,a_n]\) 的张量。
indices 是一个形状为 \([i_0,...,i_m]\) 的张量。 当 mode 设置为 ELEMENT 时， \(m=n\)。
output 是一个形状为
当
mode设置为DEFAULT时： \([a_0,..., a_{axis-1}, i_{\text{num_elementwise_dims}},..., i_m, a_{axis+1},...,a_n]\)， 并且它的秩是 n + m - 1 - num_elementwise_dims。当
mode设置为ELEMENTS时： \([i_0,...,i_m]\)。当
mode设置为ND时： 张量的秩是 \(n + m - i_m - 1 - axis\)。
， 其中 0 < num_elementwise_dims <= 1。
示例¶
Gather
in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(1, 3))
layer = network.add_gather(in1, indices, axis=2)
network.mark_output(layer.get_output(0))
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]]])
inputs[indices.name] = np.array([2, 1, 1])
outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[[-1.0, -2.0, -2.0]]], [[[2.0, 1.0, 1.0]]]])
使用 2D 索引的 Gather
in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 2))
layer = network.add_gather(in1, indices, axis=2)
network.mark_output(layer.get_output(0))
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]]])
inputs[indices.name] = np.array([[2, 1], [3, 3]])
outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[[-1.0, -2.0], [10.0, 10.0]]], [[[2.0, 1.0], [-2.0, -2.0]]]])
GatherElements
in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 1, 5))
layer = network.add_gather_v2(in1, indices, mode=trt.GatherMode.ELEMENT)
layer.axis = 2
network.mark_output(layer.get_output(0))
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]]])
inputs[indices.name] = np.array([[[0, 1, 2, 3, 4]], [[4, 2, 3, 3, 1]]])
outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[-3.0, -2.0, -1.0, 10.0, -25.0]], [[-1.0, 2.0, -2.0, -2.0, 1.0]]])
轴=0 的 GatherElements
in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 1, 5))
layer = network.add_gather_v2(in1, indices, mode=trt.GatherMode.ELEMENT)
layer.axis = 0
network.mark_output(layer.get_output(0))
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]]])
inputs[indices.name] = np.array([[[0, 1, 1, 0, 0]], [[1, 1, 0, 0, 1]]])
outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[-3.0, 1.0, 2.0, 10.0, -25.0]], [[0.0, 1.0, -1.0, 10.0, -1.0]]])
GatherND
in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 3))
layer = network.add_gather_v2(in1, indices, mode=trt.GatherMode.ND)
network.mark_output(layer.get_output(0))
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]]])
inputs[indices.name] = np.array([[0, 0, 2], [1, 0, 1]])
outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[-1.0, 1.0]])
C++ API¶
有关 C++ IGatherLayer 算子的更多信息，请参阅 C++ IGatherLayer 文档。
Python API¶
有关 Python IGatherLayer 算子的更多信息，请参阅 Python IGatherLayer 文档。