pytorch 等效项 tf.gather

Question

我在将一些代码从tensorflow移植到pytorch时遇到一些问题。

所以我有一个尺寸为 10x30 的矩阵，代表 10 个示例，每个示例有 30 个特征。然后我有另一个维度为 10x5 的矩阵，其中包含第一个矩阵中每个示例的 5 个最接近示例的索引。我想使用第二个矩阵中包含的索引“收集”第一个矩阵中每个示例的 5 个最接近的示例，从而得到形状为 10x5x30 的 3d 张量。

在张量流中，这是通过 完成的tf.gather(matrix1, matrix2)。有谁知道我如何在 pytorch 中做到这一点？

Answer 1

这个怎么样？

matrix1 = torch.randn(10, 30)
matrix2 = torch.randint(high=10, size=(10, 5))
gathered = matrix1[matrix2]

它使用整数数组索引的技巧。