Scatter¶

通过从输入张量复制值，然后根据给定的索引和更新张量来更新值，从而创建一个输出张量。

属性¶

mode 要使用的 scatter 模式。可以是以下之一

ELEMENT indices 和 tensors 定义了对 input 元素的更新操作。
ND indices 和 tensors 定义了对 input 切片的更新操作。

axis 要在其上进行 scatter 操作的轴。默认为 0。

输入¶

input: 类型为 T1 的张量。

indices: 类型为 T2 的张量。

updates: 类型为 T1 的张量。

输出¶

output: 类型为 T1 的张量。

数据类型¶

T1: int8, int32, int64, float16, float32, bfloat16

T2: int32, int64

形状信息¶

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

indices 是一个形状为 \([i_0,...,i_m]\) 的张量。当 mode 设置为 ELEMENT 时，\(m=n\)。

updates 是一个形状为 \([j_0,...,j_k]\) 的张量。当 mode 设置为 ELEMENT 时，\(k=n\)。当 mode 设置为 ND 时，\(k = n + m - i_m - 1 - 1\)。

output 是一个形状为 \([a_0,...,a_n]\) 的张量。

容量限制¶

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

示例¶

Scatter

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 1, 5))
updates = network.add_input("updates", dtype=trt.float32, shape=(2, 1, 5))
layer = network.add_scatter(in1, indices, updates, mode=trt.ScatterMode.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([[[1, 2, 3, 0, 4]], [[4, 1, 2, 3, 0]]])

inputs[updates.name] = np.array([[[-1.0, 2.4, 3.2, 10.8, 8.9]], [[0, -11.2, 34.2, 223.9, -100]]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[10.8, -1.0, 2.4, 3.2, 8.9]], [[-100, -11.2, 34.2, 223.9, 0]]])

在轴 0 上进行 Scatter 操作

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 1, 5))
updates = network.add_input("updates", dtype=trt.float32, shape=(2, 1, 5))
layer = network.add_scatter(in1, indices, updates, mode=trt.ScatterMode.ELEMENT)
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([[[1, 0, 0, 0, 1]], [[0, 1, 1, 1, 0]]])

inputs[updates.name] = np.array([[[-1.0, 2.4, 3.2, 10.8, 8.9]], [[0, -11.2, 34.2, 223.9, -100]]])

outputs[layer.get_output(0).name] = layer.get_output(0).shape
expected[layer.get_output(0).name] = np.array([[[0, 2.4, 3.2, 10.8, -100]], [[-1.0, -11.2, 34.2, 223.9, 8.9]]])

ScatterND

in1 = network.add_input("input1", dtype=trt.float32, shape=(2, 1, 5))
indices = network.add_input("indices", dtype=trt.int32, shape=(2, 1, 3))
updates = network.add_input("updates", dtype=trt.float32, shape=(2, 1))
layer = network.add_scatter(in1, indices, updates, mode=trt.ScatterMode.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, 1]]], [[[1, 0, 3]]]])

inputs[updates.name] = np.array([[-1.0, 0]])

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

C++ API¶

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

Python API¶

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