Tensorflow pytorch等效tf.gather

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

所以我有一个尺寸为10x30的矩阵，代表10个例子，每个例子有30个特征。然后我有另一个维度为10x5的矩阵，包含第一个矩阵中每个例子的5个最接近的例子的索引。我想使用第二个矩阵中包含的索引“收集”第一个矩阵中每个示例的5个壁橱示例，给我留下一个10x5x30形状的三维张量

在tensorflow中，这是通过

tf.gather（matrix1，matrix2）

完成的。有人知道我如何在pytorch中做到这一点吗？

这个怎么样

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

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

这个怎么样

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

---

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

我不确定它在TF中是如何完成的，但是你检查过了吗？我不确定它在TF中是如何完成的，但是你检查过了吗？它会产生形状为（10，5，30）的张量。。。。它真的对OP.有效吗？它产生了形状为（10，5，30）的张量。。。。这真的对OP有效吗。。。？