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

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

1 个答案: