如何快速计算两组高维嵌入数组的全量余弦相似度？

CPU端Numpy优化方案 #

你的当前方案效率低下的核心问题是：生成笛卡尔积的DataFrame完全冗余，且np.vectorize本质是Python循环，没有利用numpy的向量化并行能力。正确的做法是利用矩阵运算直接计算全量余弦相似度，步骤如下：

原理 #

余弦相似度公式可转化为：

cos(a,b) = (a · b) / (||a|| * ||b||)

若先将所有向量归一化（即除以自身的L2范数），则余弦相似度等价于归一化向量的点积。因此只需两步：

1. 对数组A、B的所有向量分别做L2归一化
2. 计算归一化后A与B的转置矩阵的点积，得到的矩阵每个元素对应一对向量的余弦相似度

代码实现 #

import numpy as np

# 假设A是形状为(64000, 1024)的numpy数组，B是(12000, 1024)的numpy数组
# 1. 归一化向量（L2范数）
A_norm = A / np.linalg.norm(A, axis=1, keepdims=True)
B_norm = B / np.linalg.norm(B, axis=1, keepdims=True)

# 2. 计算全量余弦相似度矩阵（形状为(64000, 12000)）
cos_sim_matrix = A_norm @ B_norm.T

内存优化（分块计算） #

如果内存不足以存储完整的cos_sim_matrix（约3GB，按float32计算），可以将A分块处理：

chunk_size = 6400  # 每块处理6400个A的向量，可根据内存调整
cos_sim_chunks = []

for i in range(0, A.shape[0], chunk_size):
    A_chunk = A[i:i+chunk_size]
    A_chunk_norm = A_chunk / np.linalg.norm(A_chunk, axis=1, keepdims=True)
    chunk_sim = A_chunk_norm @ B_norm.T
    cos_sim_chunks.append(chunk_sim)
    # 可直接将chunk_sim写入文件，避免内存累积
    # np.save(f"cos_sim_chunk_{i//chunk_size}.npy", chunk_sim)

# 若需要合并所有块（可选）
cos_sim_matrix = np.concatenate(cos_sim_chunks, axis=0)

这种方式的速度会比你当前的方案提升几个数量级，CPU上每秒可处理数亿次相似度计算。

GPU端实现方案 #

该操作非常适合在GPU上运行，因为矩阵乘法和归一化都是GPU擅长的并行计算任务，实现成本很低，常用的库有PyTorch、TensorFlow或CuPy（与Numpy API兼容）。

方案1：使用PyTorch #

import torch

# 将数组移至GPU
A_tensor = torch.tensor(A, dtype=torch.float32).cuda()
B_tensor = torch.tensor(B, dtype=torch.float32).cuda()

# 归一化
A_norm = A_tensor / torch.norm(A_tensor, dim=1, keepdim=True)
B_norm = B_tensor / torch.norm(B_tensor, dim=1, keepdim=True)

# 计算相似度矩阵
cos_sim_matrix = torch.matmul(A_norm, B_norm.T)

# 若需要转回Numpy数组
cos_sim_matrix_np = cos_sim_matrix.cpu().numpy()

方案2：使用CuPy（与Numpy代码几乎一致） #

import cupy as cp

# 将数组转为CuPy数组（自动在GPU上存储）
A_cp = cp.array(A)
B_cp = cp.array(B)

# 归一化
A_norm_cp = A_cp / cp.linalg.norm(A_cp, axis=1, keepdims=True)
B_norm_cp = B_cp / cp.linalg.norm(B_cp, axis=1, keepdims=True)

# 计算相似度矩阵
cos_sim_matrix_cp = A_norm_cp @ B_norm_cp.T

# 转回Numpy数组
cos_sim_matrix_np = cp.asnumpy(cos_sim_matrix_cp)

GPU的计算速度会比CPU快10~100倍，具体取决于GPU型号。

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