pytorch中的向量和矩阵行余弦相似度

Question

在pytorch中，我有多个（十万级）300个dim向量（我想我应该在矩阵中上传），我想根据它们与另一个向量的余弦相似度对它们进行排序并提取前1000个.我想避免 for 循环，因为它很耗时。我一直在寻找一个有效的解决方案。

Answer 1

您可以使用torch.nn.functional.cosine_similarity function for computing cosine similarity. And torch.argsort提取前1000名。

这是一个例子：

x = torch.rand(10000,300)
y = torch.rand(1,300)
dist = F.cosine_similarity(x,y)
index_sorted = torch.argsort(dist)
top_1000 = index_sorted[:1000]

请注意y的形状，不要忘记在调用相似函数之前重新整形。另请注意 argsort 只是 return 最接近向量的索引。要访问这些向量本身，只需编写 x[top_1000]，这将 return 一个形状为 (1000,300) 的矩阵。

A vector and matrix rows cosine similarity in pytorch