别再被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处理），行为会变得复杂：

python复制a = torch.randn(3, 4, 5)  # batch_size=3, seq_len=4, dim=5
b = torch.randn(3, 4, 5)

提示：在Transformer等模型中处理注意力分数时，通常使用dim=-1确保在特征维度计算

3. 两两相似度矩阵计算技巧

实际项目中最常见的需求是计算两组向量间的全连接相似度。假设：

• 矩阵A形状为(m,d)
• 矩阵B形状为(n,d)
• 需要得到(m,n)的相似度矩阵

3.1 广播机制解法

python复制def pairwise_cosine_sim(A, B):
    # 扩展维度：A (m,1,d) & B (1,n,d)
    A = A.unsqueeze(1)  # shape: (m,1,d)
    B = B.unsqueeze(0)  # shape: (1,n,d)
    
    # 广播计算：(m,n,d) -> (m,n)
    return F.cosine_similarity(A, B, dim=-1)

原理拆解：

1. 通过unsqueeze引入广播维度
2. 自动扩展后实际比较的是A的每个(m,1,d)与B的每个(1,n,d)
3. dim=-1指定在最后一个维度（特征维度）计算点积

3.2 内存优化版本

当处理大规模数据时，可改用矩阵运算实现：

python复制def pairwise_cosine_sim_mem(A, B):
    A_norm = A / A.norm(dim=1, keepdim=True)
    B_norm = B / B.norm(dim=1, keepdim=True)
    return torch.mm(A_norm, B_norm.T)  # (m,d) @ (d,n) -> (m,n)

两种方法的性能对比（RTX 3090测试）：

4. 工程实践中的常见陷阱

4.1 维度不匹配错误

错误示例：

python复制a = torch.randn(3,4)
b = torch.randn(4,5)  # 维度不一致
F.cosine_similarity(a, b)  # 报错

解决方案：

python复制# 方案1：转置对齐
b = b.T  # (5,4)
sim = F.cosine_similarity(a, b, dim=1)

# 方案2：广播计算
a = a.unsqueeze(1)  # (3,1,4)
b = b.unsqueeze(0)  # (1,5,4)
sim = F.cosine_similarity(a, b, dim=-1)  # (3,5)

4.2 数值稳定性问题

当向量接近零向量时会出现除零错误：

python复制zero_vec = torch.zeros(10)
F.cosine_similarity(zero_vec, zero_vec)  # 输出nan

改进方案：

python复制def safe_cosine_sim(a, b, eps=1e-8):
    dot = (a * b).sum(dim=-1)
    norm = a.norm(dim=-1) * b.norm(dim=-1) + eps
    return dot / norm

4.3 混合精度训练中的问题

在AMP自动混合精度下，可能出现精度损失：

python复制with torch.cuda.amp.autocast():
    # 可能得到不稳定的结果
    sim = F.cosine_similarity(half_a, half_b)  

解决方案：

python复制with torch.cuda.amp.autocast(enabled=False):
    # 强制使用FP32计算
    sim = F.cosine_similarity(half_a.float(), half_b.float())

5. 典型应用场景实现

5.1 文本相似度计算

python复制from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")

def text_similarity(text1, text2):
    inputs = tokenizer([text1, text2], return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1)  # (2,768)
    return F.cosine_similarity(embeddings[0], embeddings[1], dim=0)

5.2 图像检索系统

python复制from torchvision.models import resnet50
from torchvision import transforms

model = resnet50(pretrained=True).eval()
preprocess = transforms.Compose([...])

def image_search(query_img, gallery_imgs):
    # query_img: PIL Image
    # gallery_imgs: List[PIL Image]
    query_feat = model(preprocess(query_img).unsqueeze(0))  # (1,2048)
    gallery_feats = torch.cat([model(preprocess(img).unsqueeze(0)) for img in gallery_imgs])
    return pairwise_cosine_sim(query_feat, gallery_feats).squeeze(0)

5.3 推荐系统中的用户-物品匹配

python复制user_emb = torch.randn(1000, 128)  # 用户嵌入
item_emb = torch.randn(5000, 128)  # 物品嵌入

def recommend(user_idx, top_k=10):
    sim_scores = pairwise_cosine_sim(user_emb[user_idx:user_idx+1], item_emb)
    return torch.topk(sim_scores, k=top_k)

6. 性能优化技巧

6.1 半精度加速

python复制def optimized_pairwise_sim(A, B):
    A = A.half()  # FP16
    B = B.half()
    A = A / (A.norm(dim=-1, keepdim=True) + 1e-6)
    B = B / (B.norm(dim=-1, keepdim=True) + 1e-6)
    return torch.matmul(A, B.T).float()  # 转回FP32避免累积误差

6.2 分块计算策略

处理超大规模矩阵时（如10万x10万）：

python复制def chunked_cosine_sim(A, B, chunk_size=5000):
    sim_matrix = []
    for i in range(0, len(A), chunk_size):
        chunk_sim = []
        for j in range(0, len(B), chunk_size):
            chunk = pairwise_cosine_sim(A[i:i+chunk_size], B[j:j+chunk_size])
            chunk_sim.append(chunk)
        sim_matrix.append(torch.cat(chunk_sim, dim=1))
    return torch.cat(sim_matrix, dim=0)

6.3 GPU内存优化

通过梯度检查点减少显存占用：

python复制from torch.utils.checkpoint import checkpoint

class CosineSimWithCheckpoint(torch.nn.Module):
    def forward(self, A, B):
        return checkpoint(pairwise_cosine_sim, A, B)

7. 与其他相似度计算的对比

在PyTorch中的实现对比：

python复制# 余弦相似度
sim = F.cosine_similarity(a, b)

# 欧式距离
dist = torch.cdist(a, b, p=2)

# 点积相似度
sim = torch.matmul(a, b.T)

# 曼哈顿距离
dist = torch.cdist(a, b, p=1)