在给定索引数组和 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来做，但如果这一切都可以用火炬完成，那就太好了

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])