沿特定维度和特定通道索引整个张量

【问题标题】：Indexing whole Tensor along specific dimension and specific channels沿特定维度和特定通道索引整个张量
【发布时间】：2021-07-17 17:49:20
【问题描述】：

假设我们有一个张量 A，其维度为 dim(A)=[i, j, k=6, u, v]时间>。现在我们有兴趣用 channels=[0:3] 获得维度 k 的整个张量。我知道我们可以通过这种方式得到它：

B = A[:, :, 0:3, :, :]

现在我想知道是否有更好的“pythonic”方法来实现相同的结果，而无需进行这种次优索引。我的意思是这样的。

B = subset(A, dim=2, index=[0, 1, 2])

无论在哪个框架下，即pytorch、tensorflow、numpy等

非常感谢

【问题讨论】：

tf.gather 将不必要的复杂，在您的情况下，您必须构建索引数组，例如 tf.gather(A,tf.tile([[[0,1,2]]],[10,10,1]),batch_dims=2)

【解决方案1】：

在numpy中，可以使用take方法：

B = A.take([0,1,2], axis=2)

在 TensorFlow 中，没有比使用传统方法更简洁的方法了。使用tf.slice 会非常冗长：

B = tf.slice(A,[0,0,0,0,0],[-1,-1,3,-1,-1])

您可以使用take 的实验版本（自 TF 2.4 起）：

B = tf.experimental.numpy.take(A, [0,1,2], axis=2)

在 PyTorch 中，你可以使用index_select:

torch.index_select(A, dim=2, index=torch.tensor([0,1,2]))

请注意，您可以通过使用 ellipsis 来跳过明确列出第一个维度（或最后一个维度）：

# Both are equivalent in that case
B = A[..., 0:3, :, :]
B = A[:, :, 0:3, ...]

【讨论】：