为何SciPy的distance.cosine函数比手动Python实现更快？

我用两段代码计算相同向量的余弦相似度，手动实现基本照搬了SciPy的内部逻辑，甚至移除了一些检查，本以为会更快，但实际测试下来SciPy的调用速度反而略快（约0.55ms vs 约0.69ms），想知道原因。

测试代码 #

import time
import math

import numpy as np
from scipy.spatial import distance

SIZE = 6400000
EXECUTIONS = 10000

path = "" # 向量文件路径
file_data = np.genfromtxt(path, delimiter=',')
A,B = np.moveaxis(file_data, 1, 0).astype('f')

start_time = time.time()
for _ in range(EXECUTIONS):    
    cos_sim = distance.cosine(A,B)

print("SciPy 耗时：%s ms" % (((time.time() - start_time) * 1000)/EXECUTIONS))
cos_sim_scipy = cos_sim

def cosine(u, v, w=None):
    uv = np.dot(u, v)
    uu = np.dot(u, u)
    vv = np.dot(v, v)
    dist = 1.0 - uv / math.sqrt(uu * vv)
    # 修正 rounding error
    return np.clip(dist, 0.0, 2.0)

start_time = time.time()
for _ in range(EXECUTIONS):
    cos_sim = cosine(A,B)

print("手动实现 耗时：%s ms" % (((time.time() - start_time) * 1000)/EXECUTIONS))
cos_sim_manual = cos_sim

print(np.isclose(cos_sim_scipy, cos_sim_manual))

补充：向量生成代码 #

def generate_random_vector(size):
    """生成两个指定大小的随机向量并保存到文件"""
    A = np.random.normal(loc=1.5, size=(size,))
    B = np.random.normal(loc=-1.5, scale=2.0, size=(size,))
    vectors = np.stack([A, B], axis=1)
    np.savetxt('vectors.csv', vectors, fmt='%f,%f')

generate_random_vector(640000)

测试环境 #

• AMD Ryzen 9 3900X 12核处理器
• 64GB RAM
• Debian 12
• Python 3.11.2
• scipy 1.13.0
• numpy 1.26.4

核心原因分析 #

• SciPy是编译后的C级实现：你手动写的是Python函数，哪怕内部调用numpy的C实现，每次函数调用都要经过Python解释器的调度、参数检查等环节，累计10000次调用后，这些开销就会明显体现。而SciPy的distance.cosine是通过Cython编译成机器码的，直接和numpy的底层C API交互，几乎没有Python层面的额外开销。
• 更高效的底层计算逻辑：你手动实现里用了math.sqrt(uu * vv)，math.sqrt是Python标准库函数，调用开销比SciPy内部用的C级sqrt函数（或numpy的np.sqrt）更高。另外，SciPy计算向量范数时，可能直接调用了BLAS/LAPACK的优化接口（比如dnrm2），这类函数是专门为范数计算优化过的，能利用CPU的SIMD指令集加速，比两次np.dot再加开根号的组合效率更高。
• 内存与类型的极致优化：SciPy内部会确保输入数据的内存布局是连续的（比如C顺序），同时对数据类型做严格匹配，避免numpy在计算时做隐式的类型转换或内存拷贝。你的手动实现虽然也转成了float32，但函数调用时的参数传递可能会有额外的检查或临时对象创建，累计下来影响了速度。

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