Python嵌套for循环优化：大规模矩阵余弦相似度计算提速方案

场景说明 #

现有两个列数一致的待计算矩阵：

• x_tsvd：460万行，260列
• svd_tfidf：1862行，260列
  需求为计算x_tsvd中每一行与svd_tfidf所有行的余弦相似度，为x_tsvd每行匹配svd_tfidf中相似度最高的前6条结果。
  初始实现采用双层嵌套for循环，效率极低；移除内层循环的初步优化版本仍存在大量冗余计算，有非常大的性能提升空间。

初始双层循环代码：

from numpy.linalg import norm
best_match=[]

keys=np.array(df_5M['file'])
values=np.array(df['file'])

for i in range(len(x_tsvd)):
    array_=[]
    for j in range (len(svd_tfidf)):
        cosine_similarity_=np.dot(x_tsvd[i],svd_tfidf[j])/(norm(x_tsvd[i])*norm(svd_tfidf[j]))
        array_.append(cosine_similarity_)
    index=np.array(array_).argsort()
    best_match.append({keys[i]:values[index][::-1][0:5]})

初步优化（移除内层循环）代码：

from numpy.linalg import norm
best_match=[]
keys=np.array(df_5M['file'])
values=np.array(df['file'])
for i in range(len(x_tsvd)):
    a=x_tsvd[i]
    b=svd_tfidf
    a_dot_b=np.sum(np.multiply(a,b),axis=1)
    norm_a=norm(a)
    norm_b=norm(b,axis=1)
    cosine_similarity_=a_dot_b/(norm_a*norm_b)
    index=np.argsort(cosine_similarity_)
    best_match.append({keys[i]:values[index][::-1][0:6]})       

现有实现核心性能问题 #

• 双层循环版本总迭代次数超85亿，Python层循环开销极大，完全没有利用numpy的向量化计算能力，耗时不可接受。
• 初步优化版本仍保留460万次Python外层循环，且存在大量冗余计算：svd_tfidf的行范数是固定值，不需要每次循环重复计算；逐行切片、逐行全量排序的开销也非常高。

可落地优化方案 #

1. Numpy向量化优化（无额外依赖，速度较初步版本提升10~20倍） #

核心优化逻辑：

• 提前预计算固定不变的矩阵行范数，避免循环内重复计算
• 提前对两个矩阵做L2行归一化：余弦相似度公式为cos(a,b) = (a·b)/(||a||*||b||)，归一化后所有行向量的L2范数为1，两个矩阵的点积结果直接等于余弦相似度，省掉逐行做除法的开销
• 用矩阵乘法批量计算所有相似度对，把Python层循环降到最低
• 用argpartition替代全量argsort取TopN结果，排序复杂度从O(nlogn)降到O(n)，仅对取出的TopN小范围结果做精确排序即可
• 针对内存不足的场景，支持分块批量计算，不需要一次性加载全量相似度矩阵

实现代码：

import numpy as np

# 常量定义
TOP_N = 6
# 可根据自身内存调整分块大小，单块10万行时单精度下相似度矩阵约占700MB内存
BATCH_SIZE = 100000

# 加载元数据
keys = np.array(df_5M['file'])
values = np.array(df['file'])

# 提前归一化svd_tfidf，全程只需要计算一次
svd_norm = np.linalg.norm(svd_tfidf, axis=1, keepdims=True)
svd_normed = svd_tfidf / svd_norm
# 提前归一化x_tsvd
x_norm = np.linalg.norm(x_tsvd, axis=1, keepdims=True)
x_normed = x_tsvd / x_norm

best_match = []
# 分块计算避免内存溢出
for start_idx in range(0, len(x_normed), BATCH_SIZE):
    end_idx = min(start_idx + BATCH_SIZE, len(x_normed))
    batch_x = x_normed[start_idx:end_idx]
    # 批量计算当前批次所有行和svd_tfidf的余弦相似度
    batch_cos_sim = batch_x @ svd_normed.T
    # 先分区取出TopN的索引，不做全量排序
    batch_topn_idx = np.argpartition(batch_cos_sim, -TOP_N, axis=1)[:, -TOP_N:]
    # 仅对TopN范围内的相似度做精确排序，得到从高到低的索引
    batch_topn_sorted = np.take_along_axis(
        batch_topn_idx,
        np.take_along_axis(batch_cos_sim, batch_topn_idx, axis=1).argsort(axis=1)[:, ::-1],
        axis=1
    )
    # 组装当前批次结果
    for row_offset in range(len(batch_x)):
        global_row_idx = start_idx + row_offset
        best_match.append({keys[global_row_idx]: values[batch_topn_sorted[row_offset]]})

如果机器内存足够（单精度下相似度矩阵约占32GB，双精度约65GB），可以去掉分块逻辑，直接全量计算cos_sim = x_normed @ svd_normed.T，速度会更快。

2. GPU加速优化（有NVIDIA显卡时使用，速度较CPU版本再提升10~50倍） #

如果有可用的NVIDIA显卡，可以直接用CuPy替换Numpy，计算逻辑和上述代码完全一致，矩阵乘法会自动调用GPU核心并行计算，460万行的全量计算可以在几十秒内完成，只需要在代码开头把numpy替换为cupy即可：

import cupy as np
# 后续逻辑完全不变，注意提前把numpy数组转成cupy数组送入显存计算即可

3. 超大规模数据极致优化 #

如果后续数据规模进一步增长，可以直接用FAISS库做内积检索，FAISS内置了高度优化的近邻检索实现，支持CPU/GPU部署、支持量化压缩降低内存占用，亿级数据规模下的TopN检索性能远高于手写Numpy实现。

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