pytorch - 用不同形状的张量计算余弦距离

Question

我有以下代表词向量的张量

A = (2, 500)

其中第一个维度是BATCH维度(即A包含两个词向量，每个词向量有500个元素)

我还有下面的张量

B = (10, 500)

我想计算 A 和 B 之间的余弦距离，这样我得到

C = (2, 10, 1)

即对于 A 中的每一行计算与 B 中每一行的余弦距离

我查看了使用 torch.nn.functional.F.cosine_similarity 但这不起作用，因为尺寸必须相同。

在 pytorch 中实现这一目标的最有效方法是什么？

Answer 1

接受的解决方案似乎效率低下——它在我的机器上花了很长时间，最后由于内存不足而破坏了内核——而这个解决方案只花了几毫秒:

import torch.nn.functional as F

# cosine similarity = normalize the vectors & multiply
C = F.normalize(A) @ F.normalize(B).t()

这是 sentence-transformers 中的实现