如何使用numpy与并行处理加速相似度函数的计算效率？

你当前的计算逻辑是余弦相似度计算，针对你(8000, 127000)维度的固定矩阵M、动态向量N的实时搜索场景，可按以下优先级优化：

1 预处理优化（比自行实现并行收益更高） #

• 预计算固定值：np.linalg.norm(M, axis=1)是和N无关的固定值，提前计算后存储为向量norm_M，每次计算仅需计算动态N的范数norm_N = np.linalg.norm(N)，直接减少近一半运算量
• 替换优化版线性代数库：检查当前numpy是否链接了MKL/OpenBLAS这类优化BLAS库（conda安装的numpy默认带MKL），原生优化的点积运算已经实现了底层多线程加速，比自行写并行逻辑效率高3-10倍
• 降低数据精度：将M存储为float32类型，相似度计算精度损失可忽略，运算速度可提升1倍，内存占用也减少一半

优化后你的计算逻辑可以简化为：

# 初始化阶段预计算，仅执行一次
norm_M = np.linalg.norm(M.astype(np.float32), axis=1)
norm_M[norm_M == 0] = 1e-8 # 避免除0错误

# 每次动态请求执行的逻辑
def calc_similarity(N):
    N = N.astype(np.float32)
    dot_res = M.dot(N)
    norm_N = np.linalg.norm(N)
    return dot_res / (norm_M * norm_N)

2 适合按需调用的并行方案 #

你担心的多进程启动开销问题可以通过常驻进程池解决：

• 程序初始化阶段就启动和CPU物理核心数一致的常驻进程池，不要每次请求才启动进程，示例用concurrent.futures实现：

# 初始化阶段执行，仅执行一次
import os
# 关闭numpy自带的多线程，避免和多进程竞争资源
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
import numpy as np
from concurrent.futures import ProcessPoolExecutor

# 把M按行拆成和核心数一致的分片
worker_num = os.cpu_count()
M_splits = np.array_split(M, worker_num)
pool = ProcessPoolExecutor(max_workers=worker_num)

# 每次请求调用的并行计算逻辑
def calc_similarity_parallel(N):
    N = N.astype(np.float32)
    norm_N = np.linalg.norm(N)
    # 并行计算各分片的点积
    dot_list = list(pool.map(lambda x: x.dot(N), M_splits))
    dot_res = np.concatenate(dot_list)
    return dot_res / (norm_M * norm_N)

• 也可以用numba的JIT编译加速，给自定义计算函数加@njit(parallel=True)装饰器，第一次调用编译完成后，后续调用速度比原生numpy快20%-50%，适合实时调用场景

3 极端低延迟场景优化 #

如果需要压到百毫秒级延迟，可以直接使用成熟的向量检索工具：

• 用faiss提前为M构建Flat索引，直接调用faiss的相似度检索接口，底层做了SIMD指令和多线程极致优化，8000条127000维向量的检索延迟可压到100ms以内
• 若允许极小的精度损失，可提前用PCA将127000维向量降为256/512维，运算速度可提升上百倍

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