Python TensorFlow ``gather_nd()`` 方法详解与代码示例


发布日期 : 2023-06-18 11:32:05 UTC

访问量: 10 次浏览

Python – tensorflow.gather\_nd()

TensorFlow是谷歌设计的开源Python库,用于开发机器学习模型和深度学习神经网络。

gather\_nd()用于根据提供的指数从输入张量中收集切片。

语法:
tensorflow.gather_nd( params, indices, batch_dims, name)

参数:

  • params : 它是一个等级大于或等于 axis+1 的张量。
  • indices:它是一个 dtype int32int64 的张量。
  • batch\_dims :它是一个整数,代表批次维度的数量。它必须小于rank(indices)。
  • name : 它定义了操作的名称。

返回值:

它返回一个张量,其 dtype 与参数相同。

示例 1:

# Importing the library
import tensorflow as tf

# Initializing the input
data = tf.constant([[1, 2], [3, 4], [5, 6]])
indices = tf.constant([[1], [0], [1]])

# Printing the input
print('data: ',data)
print('indices: ',indices)

# Calculating result
res = tf.gather_nd(data, indices)

# Printing the result
print('res: ',res)

输出:

data: tf.Tensor(
[[1 2]
[3 4]
[5 6]], shape=(3, 2), dtype=int32)
indices: tf.Tensor(
[[1]
[0]
[1]], shape=(3, 1), dtype=int32)
res: tf.Tensor(
[[3 4]
[1 2]
[3 4]], shape=(3, 2), dtype=int32)

示例 2:

# Importing the library
import tensorflow as tf

# Initializing the input
data = tf.constant([[1, 2, 3], [3, 4, 5], [5, 6, 7]])
indices = tf.constant([[1, 0], [0, 2], [1, 2]])

# Printing the input
print('data: ',data)
print('indices: ',indices)

# Calculating result
res = tf.gather_nd(data, indices)

# Printing the result
print('res: ',res)

输出:

data: tf.Tensor(
[[1 2 3]
[3 4 5]
[5 6 7]], shape=(3, 3), dtype=int32)
indices: tf.Tensor(
[[1 0]
[0 2]
[1 2]], shape=(3, 2), dtype=int32)
res: tf.Tensor([3 3 5], shape=(3,), dtype=int32)