TensorFlow - 按轴索引

Question

我想索引一个任意形状的张量的最后一个轴，除了最后一个是 2。

例如令 x 的形状为 (1,2,2)。索引到最后一个轴

x_0 = x[:, :, 0]    # x_0, x_1 shapes are (1,2)
x_1 = x[:, :, 1]

例如令 x 的形状为 (1,2,3,4,2)。索引到最后一个轴

x_0 = x[:, :, :, :, 0]   # x_0, x_1 shapes are (1,2,3,4)
x_1 = x[:, :, :, :, 1]

我一直找不到任何用于切片任意形状的 tensorflow 函数或用法。

我需要一个通用的索引方法，这样我就可以随时访问任何形状张量的最后一个轴。

Answer 1

tensorflow 中的切片语法与numpy 非常相似。在这种情况下，您可以使用ellipsis：

省略号扩展为选择元组索引所有维度所需的: 对象的数量。在大多数情况下，这意味着扩展选择元组的长度是x.ndim。可能只有一个省略号。

在你的情况下，

x_0 = x[..., 0]

将索引具有任意形状的张量的最后一个轴。

你也可以看看the answer to the question: What is the difference between the slice (:) and the ellipsis (…) operators in numpy?。