PyTorch 张量：基于旧张量和索引的新张量答案

【问题标题】：PyTorch tensors: new tensor based on old tensor and indicesPyTorch 张量：基于旧张量和索引的新张量
【发布时间】：2021-06-27 07:17:14
【问题描述】：

我是张量新手，对这个问题很头疼：

我有一个大小为 k 的索引张量，其值介于 0 和 k-1 之间：

tensor([0,1,2,0])

以及以下矩阵：

tensor([[[0, 9],
     [1, 8],
     [2, 3],
     [4, 9]]])

我想创建一个新的张量，其中包含索引中指定的行，按该顺序。所以我想要：

tensor([[[0, 9],
     [1, 8],
     [2, 3],
     [0, 9]]])

在张量之外，我或多或少会这样操作：

new_matrix = [matrix[i] for i in index]

如何在 PyTorch 中对张量执行类似操作？

【问题讨论】：

【解决方案1】：

你使用fancy indexing:

from torch import tensor

index = tensor([0,1,2,0])
t = tensor([[[0, 9],
     [1, 8],
     [2, 3],
     [0, 9]]])

result = t[:, index, :]

得到

tensor([[[0, 9],
         [1, 8],
         [2, 3],
         [0, 9]]])

注意t.shape == (1, 4, 2) 和你想在 second 轴上索引；所以我们在第二个参数中应用它并通过:s 保持其余部分相同，即[:, index, :]。

【讨论】：

@SebastianT.Vincent 很高兴它成功了！如果您愿意，请考虑在时间允许的情况下接受以向其他人表明问题已解决。
一个后续问题：如果索引张量是 2 维或 n 维，我们还能使用花式索引吗？我有一个数据张量 [32,4,1000] 和一个索引张量 [32,4]，做 data[index, :] 会产生一个 [32,4,4,1000] 张量！
@SebastianT.Vincent 是的，花式索引适用于任何 n 维数组，但我不太确定您的情况应该是什么结果。您能否详细说明输出的形状应该是什么？
我希望输出的形状与输入的形状相同（即 [32,4,1000]）。我们称 dim0 批处理 (32)、dim1 光束 (4)。在每个批次中，我希望根据该批次的索引矩阵更换梁。所以在第 2 批中，如果 col[2,:] 是 [0,1,2,0]，那么如果第 2 批的梁是 ABCD，那么我希望它们被替换为 ABCA。这有意义吗？
@SebastianT.Vincent 我明白了，我认为这是一个解决方案：data[np.arange(len(data)), index.T, :].permute(1, 0, 2)。由于我们需要为每个批次应用不同的索引集，我们使用np.arange（例如，等于0..31）来索引批次维度。然后，由于梁堆叠在每批的列中，我们需要以[indices, :] 的方式选择那里。但是索引是(batch_size, beam_dim)，所以我们需要转置来对齐它们。最后一个维度是免费的，因此是:。最后我们permute第一个和第二个维度，因为我们必须做转置。