在给定索引数组和 pytorch 张量的情况下查找欧几里得距离

Question

我有一个 pytorch 张量：

Z = np.random.rand(100,2)
tZ = autograd.Variable(torch.cuda.FloatTensor(Z), requires_grad=True)

和一个索引数组：

idx = (np.array([0, 0, 0, 4, 3, 8], dtype="int64"),
       np.array([0, 1, 2, 3, 7, 4], dtype="int64"))

我需要使用 idx 数组作为索引来查找我的 tZ 张量中所有点对的距离。

现在我正在使用 numpy 来做这件事，但如果这一切都可以使用 torch 来完成就好了

dist = np.linalg.norm(tZ.cpu().data.numpy()[idx[0]]-tZ.cpu().data.numpy()[idx[1]], axis=1)

如果有人知道如何使用 pytorch 来加速它，那将是一个很大的帮助！

Answer 1

使用torch.index_select()：

Z = np.random.rand(100,2)
tZ = autograd.Variable(torch.cuda.FloatTensor(Z), requires_grad=True)

idx = (np.array([0, 0, 0, 4, 3, 8], dtype="int64"),
       np.array([0, 1, 2, 3, 7, 4], dtype="int64"))

tZ_gathered = [torch.index_select(tZ, dim=0,
                                  index=torch.cuda.LongTensor(idx[i]))
                                  # note: you may have to wrap it in a Variable too
                                  # (thanks @rvd for the comment):
                                  # index = autograd.Variable(torch.cuda.LongTensor(idx[i])))
               for i in range(len(idx))]

print(tZ_gathered[0].shape)
# > torch.Size([6, 2])

dist = torch.norm(tZ_gathered[0] - tZ_gathered[1])