PyTorch 深度解析：2026 年视角下的余弦相似度计算与工程化实践

在深度学习和自然语言处理的众多任务中，衡量两个向量或张量之间的相似度是一项基础且至关重要的工作。无论你是在构建基于 RAG（检索增强生成）的推荐系统、进行多模态文档检索，还是计算高维图像特征的距离，余弦相似度都是一个不可或缺的工具。

特别是站在 2026 年的视角，随着 AI 原生应用和 Agentic Workflows 的普及，我们不再仅仅关注模型的准确率，更关注推理的吞吐量和检索的精确度。在这篇文章中，我们将深入探讨如何使用 PyTorch 这一强大的深度学习框架来计算两个张量之间的余弦相似度。我们将从基本概念入手，通过具体的代码示例，一步步掌握 torch.nn.CosineSimilarity 的用法，并分享一些在实际开发中可能遇到的坑和性能优化技巧，以及我们如何在现代 AI 工程中应用这一技术。

什么是余弦相似度？

在开始编写代码之前，让我们先快速回顾一下余弦相似度的数学定义。简单来说，余弦相似度衡量的是两个向量在空间中的夹角。它的值范围在 -1 到 1 之间：

• 1.0：表示两个向量完全同向（最相似）。
• 0.0：表示两个向量正交（无关）。
• -1.0：表示两个向量完全反向（最不相似）。

与欧几里得距离不同，余弦相似度关注的是方向而非大小。这意味着，即使两个向量的数值量级差异很大（比如一个是 [1, 1]，一个是 [100, 100]），只要它们方向一致，余弦相似度依然认为是 1。这在处理文本 TF-IDF 特征、词向量或现代 LLM 的 Embedding 时非常有用。

PyTorch 中的核心工具：CosineSimilarity

PyTorch 为我们封装了一个非常方便的模块——torch.nn.CosineSimilarity。这个模块允许我们在指定的维度上高效地计算相似度。

语法与参数解析

我们可以通过以下方式调用这个模块：

torch.nn.CosineSimilarity(dim=1, eps=1e-8)

这里有几个关键参数需要你特别注意：

• INLINECODEd10242df (维度): 这是最重要的参数。它决定了沿着张量的哪个维度进行计算。理解这个参数对于处理多维数据（比如批量的矩阵数据）至关重要。默认值通常是 INLINECODEb227d47b，这对应于大多数情况下的特征维度。
• INLINECODEa64db1df (极小值): 这是一个数值稳定性的保护参数。在实际计算中，如果两个向量的模长（范数）非常小或接近于零，直接计算可能会导致除以零的错误（NaN）。INLINECODE63f96321 会在分母中加入一个极小值，确保计算的稳定性。

准备工作

为了确保我们的代码能够顺利运行，请确保你已经安装了 PyTorch。如果还没有，你可以使用 pip 快速安装。在我们的所有示例中，我们首先会导入库：

import torch
import torch.nn as nn

实战演练：代码示例详解

为了让你全面掌握这一工具，我们将从简单到复杂，通过几个不同的场景来演示如何计算余弦相似度。

示例 1：基础的一维向量计算

让我们从最基础的情况开始。假设我们有两个一维张量，代表两个简单的数据点或特征向量。我们需要计算它们之间的相似度。

对于一维向量，INLINECODEa00388ee 参数只能设置为 INLINECODE8573a6b3，因为这是唯一的维度。

# 导入 PyTorch 库
import torch
import torch.nn as nn

# 定义两个一维张量
# 这里我们模拟了两个具有4个特征的向量
tensor_1 = torch.tensor([0.5, 0.3, 1.2, 0.33])
tensor_2 = torch.tensor([0.3, 0.2, 1.3, 1.4])

# 打印输入向量，方便核对
print("第一个张量:", tensor_1)
print("第二个张量:", tensor_2)

# 初始化 CosineSimilarity 模块
# 对于一维向量，dim 必须设置为 0
cosine_sim = nn.CosineSimilarity(dim=0)

# 计算相似度
result = cosine_sim(tensor_1, tensor_2)

# 输出结果
print("计算出的余弦相似度:", result)

# 验证：我们可以手动计算一下公式来验证结果
# dot product (x . y)
dot_product = (tensor_1 * tensor_2).sum()
# norms ||x|| * ||y||
norm_1 = torch.norm(tensor_1)
norm_2 = torch.norm(tensor_2)
manual_result = dot_product / (norm_1 * norm_2)
print("手动验证结果:", manual_result.item())

在这个例子中，我们不仅使用了内置函数，还通过手动计算点积和范数的方式进行了验证。这样可以帮助你深入理解其背后的数学原理。

示例 2：处理二维张量与批量数据

在实际的深度学习应用中，我们很少只处理单个向量，更多的是处理“批量”数据。例如，我们可能有一个形状为 INLINECODE374b5b56 的张量。这时，INLINECODE6010921e 参数的选择就变得非常关键。

让我们看一个具体的例子，包含多行数据的矩阵。

import torch
import torch.nn as nn

# 定义两个二维张量 (模拟批量数据)
# 假设每一行是一个样本，有4个特征
tens_1 = torch.tensor([[0.2245, 0.2959, 0.3597, 0.6766],
                       [-2.2268, 0.6469, 0.3765, 0.7898],
                       [0.4577, 0.3228, 0.4699, 0.2389]])

tens_2 = torch.tensor([[0.2423, 0.4667, 0.4434, 0.3598],
                       [-0.6679, 0.6932, 0.5387, 0.2245],
                       [0.8277, 0.2597, 0.9834, 0.9987]])

print("第一个张量形状:", tens_1.shape)
print("第二个张量形状:", tens_2.shape)

# 场景 1: 在 dim=0 (列维度) 上计算
# 这通常用于计算特征列之间的相关性，或者比较不同样本在同一特征位置的关系
cos_sim_0 = nn.CosineSimilarity(dim=0)
output_0 = cos_sim_0(tens_1, tens_2)

print("
在 dim=0 上计算的结果 (列向量相似度):")
print(output_0)
print("注意：结果是形状为", output_0.shape, "的一维向量，对应于特征维度")

# 场景 2: 在 dim=1 (行维度) 上计算
# 这是最常见的场景，计算每一对样本之间的相似度
cos_sim_1 = nn.CosineSimilarity(dim=1)
output_1 = cos_sim_1(tens_1, tens_2)

print("
在 dim=1 上计算的结果 (样本相似度):")
print(output_1)
print("注意：结果是形状为", output_1.shape, "的一维向量，对应于批量维度")

示例 3：应用场景 —— 图像相似度检索

让我们看一个更贴近实际的例子。假设我们有一组图片的特征向量（通常来自卷积神经网络 CNN 的输出层），我们想找出哪些图片最相似。

import torch
import torch.nn as nn

def image_similarity_example():
    # 模拟场景：
    # query_img: 我们想要搜索的目标图片特征 (Batch=1, Features=128)
    # database_imgs: 数据库中的图片特征 (Batch=5, Features=128)
    
    query_img = torch.randn(1, 128) # 目标图片特征
    database_imgs = torch.randn(5, 128) # 数据库中的5张图片特征

    # 为了使用 CosineSimilarity，我们需要两个张量的形状可以进行广播或匹配
    # 如果我们想计算 query_img 和 database_imgs 中每一张的相似度：
    
    # 方法 A: 扩展 query_img 以匹配 database_imgs 的 batch size
    # 将 (1, 128) 扩展为 (5, 128)
    query_expanded = query_img.expand_as(database_imgs)
    
    cos = nn.CosineSimilarity(dim=1)
    similarities = cos(query_expanded, database_imgs)
    
    print("检索图片与数据库中各图片的相似度得分:")
    print(similarities)
    
    # 找出最相似的那一张
    best_match_index = torch.argmax(similarities)
    print(f"最匹配的图片索引是: {best_match_index.item()}，得分: {similarities[best_match_index].item():.4f}")

image_similarity_example()

2026 视角下的工程化实践：生产级代码与性能优化

在我们最近构建的一个企业级 RAG 系统中，我们需要处理百万级的文档向量检索。这时候，简单的循环计算或者错误的维度操作会导致严重的性能瓶颈。让我们来看看如何编写更加健壮、高效的代码。

矩阵运算的极致优化：广播机制

在之前的示例中，我们使用了 expand，但在处理超大规模数据时，频繁的内存复制是不可接受的。我们可以利用 PyTorch 的广播机制来避免显式扩展，从而节省显存。

import torch
import torch.nn as nn

def efficient_pairwise_similarity(query, keys):
    """
    计算查询向量和键向量集合之间的相似度，不进行显式内存复制。
    
    参数:
        query: (M, D)
        keys: (N, D)
    返回:
        sim: (M, N)
    """
    # 维度调整：query -> (M, 1, D), keys -> (1, N, D)
    # 这种操作在内存上非常廉价，因为它只修改了元数据
    query_expanded = query.unsqueeze(1) 
    keys_expanded = keys.unsqueeze(0)
    
    # 在最后一维 (D) 上计算余弦相似度
    cos = nn.CosineSimilarity(dim=2)
    return cos(query_expanded, keys_expanded)

# 模拟大规模数据
batch_size = 100
db_size = 10000
dim = 768 # 比如 BERT-base 的 hidden size

queries = torch.randn(batch_size, dim)
database = torch.randn(db_size, dim)

# 使用高效计算
sim_matrix = efficient_pairwise_similarity(queries, database)
print(f"相似度矩阵形状: {sim_matrix.shape}") # 应为 (100, 10000)

混合精度训练与 FP16 推理

为了在 2026 年的现代硬件（如 NVIDIA H100 或 AMD MI300）上获得最佳性能，我们强烈建议使用混合精度，尤其是在处理高维 Embedding 时。这不仅能加速计算，还能显著减少显存占用。

# 启用自动混合精度
from torch.cuda.amp import autocast

# 假设我们在 GPU 上运行
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
query_gpu = queries.to(device)
db_gpu = database.to(device)

with autocast():
    # PyTorch 会自动在适当的时机将计算转换为 FP16
    sim_matrix_amp = efficient_pairwise_similarity(query_gpu, db_gpu)

调试与常见陷阱：我们踩过的坑

在我们与 AI 结对编程的过程中，我们发现余弦相似度计算中最隐蔽的 Bug 通常源于未归一化的输入。

陷阱：NaN 问题

如果你的输入向量中包含全零的向量（例如，某些被 Padding 的文本片段），分母为零会导致 NaN。虽然 eps 参数有帮助，但在更复杂的网络结构中，梯度可能会反向传播爆炸。

最佳实践：

我们建议在传入 CosineSimilarity 之前，先使用 F.normalize 对输入进行 Layer Normalization。这不仅能保证数值稳定性，还能将余弦相似度计算转化为简单的点积。

import torch.nn.functional as F

# 在进入相似度计算前，强制归一化
# 这会将向量长度缩放到 1，此时 CosineSim(x, y) == Dot(x, y)
normalized_query = F.normalize(queries, p=2, dim=1)
normalized_db = F.normalize(database, p=2, dim=1)

# 现在计算变得极其简单且稳定
# (M, 1, D) @ (D, N) -> (M, N)  使用矩阵乘法甚至比 CosineSimilarity 模块更快！
sim_matrix_simple = torch.mm(normalized_query, normalized_db.T)

总结

在这篇文章中，我们深入探讨了如何在 PyTorch 中使用 torch.nn.CosineSimilarity 来计算张量之间的余弦相似度。我们从基本的数学概念出发，介绍了核心 API 的参数，并通过不同的代码示例（一维向量、二维矩阵、图像检索场景）覆盖了从入门到进阶的各种用法。

更重要的是，我们结合了 2026 年的技术趋势，分享了在生产环境中如何利用广播机制和归一化技巧来优化性能，以及如何避免常见的数值稳定性陷阱。掌握这一工具，你将能够在构建分类器、检索系统或损失函数时更加得心应手。希望这些示例和解释能帮助你更好地理解和应用这一技术。

下次当你需要衡量两个特征有多相似时，你就知道该怎么做了！如果你还有关于特定应用场景的疑问，不妨尝试自己动手写一段代码，或者深入研究 PyTorch 的官方文档获取更多细节。祝你编码愉快！