PyTorch 中的批处理余弦相似度

Question

Inputs:

• Tensor a of shape [batch_size, n, d]
• Tensor b of shape [batch_size, m, d]

Output:

• Tensor c of shape [batch_size, n, m] where c[i, j, k] is the cosine similarity between a[i, j] and b[i, k]

How to implement this efficiently in PyTorch (preferably without for loops)?

Answer 1

try this:

c = torch.cosine_similarity(a.unsqueeze(2), b.unsqueeze(1), dim=-1)