别再被dim参数搞晕了！PyTorch F.cosine_similarity实战避坑指南（附两两相似度计算）

在自然语言处理、推荐系统或图像检索项目中，计算向量相似度是基础但关键的操作。PyTorch提供的F.cosine_similarity函数看似简单，但dim参数的灵活性和广播机制的隐式规则常常让开发者陷入调试困境。本文将用可复现的代码示例，带你穿透维度迷雾，掌握从基础计算到高级用法的完整技能树。

1. 理解余弦相似度的核心逻辑

余弦相似度衡量的是两个向量在方向上的差异，与向量长度无关。数学定义为：

code复制cos(θ) = (A·B) / (||A|| * ||B||)

在PyTorch中实现时，需要特别注意三个要点：

• 输入张量必须维度相同：如果形状为(3,4)和(4,)，需要先unsqueeze
• 计算结果范围：输出值域[-1,1]，1表示完全相同，-1表示完全相反
• 批量处理能力：函数原生支持批量计算，无需手动循环

典型误区：很多开发者误以为该函数会自动进行归一化处理。实际上输入的向量如果未经L2归一化，计算结果可能不符合预期。

2. dim参数的行为解密

通过对比实验揭示不同dim设置的实际效果：

2.1 二维张量场景

给定测试数据：

python复制import torch.nn.functional as F
a = torch.tensor([[1., 2], [3, 4]])  # shape (2,2)
b = torch.tensor([[5., 6], [7, 8]])  # shape (2,2)

Case 1: dim=0 (按列计算)

python复制res = F.cosine_similarity(a, b, dim=0)
# 等价于计算：
# [cos_sim([1,3], [5,7]), cos_sim([2,4], [6,8])]
# 输出：tensor([0.9558, 0.9839])

Case 2: dim=1 (按行计算)

python复制res = F.cosine_similarity(a, b, dim=1)  
# 等价于计算：
# [cos_sim([1,2], [5,6]), cos_sim([3,4], [7,8])]
# 输出：tensor([0.9734, 0.9972])

关键发现：

• dim参数决定沿着哪个维度进行向量提取
• 默认dim=1（PyTorch 1.7+版本）
• 对于二维输入，dim=0按列、dim=1按行的规律成立

2.2 高维张量场景

当处理三维张量时（如batch处理），行为会变得复杂：

``