Numpy.linalg.norm性能为何未随数据维度按比例缩放？

初始代码与性能异常 #

我编写了一段K-means聚类的子例程代码，用于将每个点分配给最近的质心，代码如下：

import numpy as np

n = 20000
D = 30
K = 250

points = np.random.rand(n, D)
centroids = np.random.rand(K, D)
membership = np.zeros(shape=n, dtype=int)
for i in range(n):
    distances = np.apply_along_axis(lambda x: np.linalg.norm(x, ord=2), 1, centroids - points[i])
    membership[i] = np.argmin(distances)

该代码的理论时间复杂度为O(NKD)，其中D是数据点的维度，按预期运行时间应随D的增减成比例变化，但实际测试结果显示运行时间随D变化很小：

D = 1
python3 benchmark.py  12.10s user 0.39s system 118% cpu 10.564 total

D = 30
python3 benchmark.py  12.17s user 0.36s system 117% cpu 10.703 total

D = 300
python3 benchmark.py  13.30s user 0.31s system 115% cpu 11.784 total

D = 1000
python3 benchmark.py  16.51s user 1.76s system 110% cpu 16.524 total

请问我忽略了什么因素？

优化后的性能提升 #

根据@Warren的建议，我修改代码直接使用np.linalg.norm的axis参数，优化后的核心代码如下：

import numpy as np

n = 20000
D = 30
K = 250

points = np.random.rand(n, D)
centroids = np.random.rand(K, D)
membership = np.zeros(shape=n, dtype=int)
for i in range(n):
    distances = np.linalg.norm(centroids - points[i], axis=1)
    membership[i] = np.argmin(distances)

性能测试结果如下，运行时间随D的变化开始符合预期：

D = 1
python3 benchmark.py  1.45s user 0.37s system 634% cpu 0.287 total

D = 30
python3 benchmark.py  1.67s user 0.29s system 592% cpu 0.331 total

D = 300
python3 benchmark.py  3.03s user 0.32s system 234% cpu 1.428 total

D = 1000
python3 benchmark.py  6.32s user 2.73s system 126% cpu 7.177 total

问题解析 #

1. Python循环与函数调用开销主导初始代码运行时间
   初始代码中，np.apply_along_axis本质是在Python层面循环调用lambda函数，每次迭代都带来额外的函数调用和循环开销。当D较小时，这些开销占总运行时间的比例极高，掩盖了D变化带来的向量计算量差异；只有当D增大到一定程度（如1000），向量计算的时间才开始超过这些开销，导致总运行时间明显上升。
2. 未充分利用numpy的向量化优势
   初始代码的外层for循环遍历每个点，加上apply_along_axis的内部Python循环，双重解释器层面的循环开销远大于实际的向量运算，使得D的变化对总时间的影响被稀释。
3. 优化代码消除冗余开销
   直接使用np.linalg.norm(..., axis=1)后，计算逻辑完全由numpy的C级向量化实现，消除了Python层面的循环和函数调用开销。此时运行时间的主导因素变为实际的向量运算量，因此随D增大，运行时间呈现出符合O(NKD)复杂度的增长趋势。

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