index_select

paddle.fluid.layers.index_select(input, index, dim=0)[源代码]

该OP沿着指定维度 dim 对输入 input 进行索引，取 index 中指定的相应项，然后返回到一个新的张量。这里 index 是一个 1-D 张量。除 dim 维外，返回的张量其余维度大小同输入 input ， dim 维大小等于 index 的大小。

参数

  • input （Variable）– 输入张量。
  • index （Variable）– 包含索引下标的一维张量。
  • dim (int, optional) – 索引轴，若未指定，则默认选取第一维。

返回

-Variable ，数据类型同输入。

代码示例

  1. import paddle
  2. import paddle.fluid as fluid
  3. import numpy as np
  5. data = np.array([[1.0, 2.0, 3.0, 4.0],
  6.                     [5.0, 6.0, 7.0, 8.0],
  7.                     [9.0, 10.0, 11.0, 12.0]])
  8. data_index = np.array([0, 1, 1]).astype('int32')
  10. with fluid.dygraph.guard():
  11.     x = fluid.dygraph.to_variable(data)
  12.     index = fluid.dygraph.to_variable(data_index)
  13.     out_z1 = fluid.layers.index_select(x, index)
  14.     print(out_z1.numpy())
  15.     #[[1. 2. 3. 4.]
  16.     # [5. 6. 7. 8.]
  17.     # [5. 6. 7. 8.]]
  18.     out_z2 = fluid.layers.index_select(x, index, dim=1)
  19.     print(out_z2.numpy())
  20.     #[[ 1.  2.  2.]
  21.     # [ 5.  6.  6.]
  22.     # [ 9. 10. 10.]]