时间戳

首页 > 解决方案 > 无论维度如何，如何从张量中选择选定的行？

python - 无论维度如何，如何从张量中选择选定的行？

问题描述

我有一个张tensorflow.Tensor量（AAA[i,:,:,...,:]i

问题是我事先不知道A有多少轴。那么我该如何编写这个操作呢？

标签： pythonpython-3.xtensorflow2.0

解决方案

这正是tf.gather()它的用途。请参见下面的示例代码：

x = tf.reshape(tf.constant([1, 2, 3, 4, 5, 6, 7, 8]), [2, 2, 2])

# This is using tf.gather() on a 3D tensor.
print(tf.gather(x, [1]))

结果是：

<tf.Tensor: shape=(1, 2, 2), dtype=int32, numpy=
array([[[5, 6],
        [7, 8]]], dtype=int32)>

x = tf.reshape(tf.constant([1, 2, 3, 4, 5, 6, 7, 8]), [2, 4])

# This is using tf.gather() on a 2D tensor.
print(tf.gather(x, [1]))

结果是：

tf.Tensor([[5 6 7 8]], shape=(1, 4), dtype=int32)

推荐阅读