关于PyTorch中沿dim=0计算余弦相似度的结果困惑

你的理解偏差在于对dim参数的定义搞反了。F.cosine_similarity中的dim指定的是要计算相似度的向量所在的维度，而非分组维度。

你的例子拆解 #

你的张量a和b形状均为(2,4)：

• a的4个列向量依次是：[1,1]、[2,9]、[3,9]、[5,5]
• b的4个列向量依次是：[7,10]、[8,-1]、[9,3]、[5,5]

当设置dim=0时，函数会沿着第0维（行维度），将两个张量对应位置的列向量配对，分别计算余弦相似度，最终得到4个结果，和你看到的输出完全匹配：

• 列向量[1,1]与[7,10]的相似度：≈0.9848
• 列向量[2,9]与[8,-1]的相似度：≈0.0942
• 列向量[3,9]与[9,3]的相似度：=0.6000
• 列向量[5,5]与[5,5]的相似度：=1.0000

实现你预期的计算逻辑 #

如果你想计算a的第一行[1,2,3,5]与b的第一行[7,8,9,5]的相似度，以及a的第二行[1,9,9,5]与b的第二行[10,-1,3,5]的相似度，需要将dim设为1（因为行向量是沿着第1维（列维度）的向量）：

import torch
import torch.nn.functional as F

a = torch.as_tensor([[1,2, 3, 5], [1,9, 9, 5]]).float()
b = torch.as_tensor([[7, 8, 9, 5], [10, -1, 3, 5]]).float()
result = F.cosine_similarity(a, b, dim=1)
# 输出结果: tensor([0.9726, 0.5893])

内容的提问来源于stack exchange，提问作者Amir Jalilifard