优化大尺寸Numpy数组余弦相似度计算：提速控内存保结果

你需要计算形状为(96341, 1000)的Numpy数组的余弦相似度，当前使用8 vCPU、32 GB内存的机器。初始代码存在内存占用过高、计算效率偏低的问题，以下是针对性优化方案，保证结果一致的同时提升速度、控制内存。

初始代码的问题分析 #

• np.dot(A, A.T)会生成(96341, 96341)的矩阵，元素数近93亿，双精度格式下内存占用超7GB，加上p1*p2的同尺寸中间矩阵，内存压力陡增；
• 重复计算np.sum(A**2, axis=1)，浪费计算资源；
• A**2的运算效率不如元素级乘法A*A。

优化方案 #

1. 基础内存与计算优化（无第三方库） #

复用范数计算、利用广播避免冗余矩阵，同时优化元素运算：

import numpy as np

def optimized_csm(A):
    # 仅计算一次所有向量的L2范数，用A*A替代A**2提升运算速度
    p_norm = np.sqrt(np.sum(A * A, axis=1))
    # 计算点积矩阵
    num = np.dot(A, A.T)
    # 广播除法，避免生成p1*p2的大中间矩阵
    result = num / (p_norm[:, np.newaxis] * p_norm)
    # 处理范数为0的情况（避免除以0得到NaN）
    result[np.isnan(result)] = 0.0
    return result

cos = optimized_csm(vec)

优化点：减少一次范数求和计算、提升元素运算速度、通过广播降低内存开销。

2. 多核并行加速 #

利用机器的8核CPU，启用BLAS多线程计算，或直接调用优化后的第三方库实现：

方案2.1：使用scikit-learn的优化实现

sklearn.metrics.pairwise.cosine_similarity内部集成了BLAS多线程优化，内存管理更高效：

from sklearn.metrics.pairwise import cosine_similarity

cos = cosine_similarity(vec)

方案2.2：手动启用Numpy多线程

如果依赖纯Numpy，可设置BLAS线程数匹配CPU核心数（需根据你的BLAS后端调整）：

import numpy as np
import os

# 针对OpenBLAS后端
os.environ['OPENBLAS_NUM_THREADS'] = '8'
# 针对Intel MKL后端（若使用）
# os.environ['MKL_NUM_THREADS'] = '8'

def optimized_csm(A):
    p_norm = np.sqrt(np.sum(A * A, axis=1))
    num = np.dot(A, A.T)
    result = num / (p_norm[:, np.newaxis] * p_norm)
    result[np.isnan(result)] = 0.0
    return result

cos = optimized_csm(vec)

3. 分块计算（极端内存紧张场景） #

将大数组拆分为小块，逐块计算并拼接结果，大幅降低单次内存占用：

import numpy as np

def chunked_csm(A, chunk_size=2000):
    n_samples = A.shape[0]
    p_norm = np.sqrt(np.sum(A * A, axis=1))
    # 预分配结果矩阵
    result = np.zeros((n_samples, n_samples), dtype=A.dtype)
    
    for i in range(0, n_samples, chunk_size):
        end = min(i + chunk_size, n_samples)
        # 计算当前块与所有向量的点积
        chunk_dot = np.dot(A[i:end], A.T)
        # 计算当前块的相似度
        chunk_sim = chunk_dot / (p_norm[i:end, None] * p_norm)
        result[i:end] = chunk_sim
        # 释放临时变量内存
        del chunk_dot, chunk_sim
    
    result[np.isnan(result)] = 0.0
    return result

cos = chunked_csm(vec)

可根据内存情况调整chunk_size，比如2000的块大小单精度下仅占约600MB临时内存。

4. 精度-内存平衡：使用单精度浮点数 #

若业务允许，将数组转为float32格式，内存占用减半，计算速度提升：

vec_float32 = vec.astype(np.float32)
cos = cosine_similarity(vec_float32)

结果一致性验证 #

所有优化方案的结果与原始代码在浮点误差范围内完全一致，可通过以下代码验证：

cos_original = csm(vec)
cos_optimized = optimized_csm(vec)  # 或其他优化方案的结果
print(np.allclose(cos_original, cos_optimized))  # 输出True则一致

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