numpy sum与bottleneck nansum性能不符问题求助

我有一段高度依赖numpy sum的代码，约60%的运行时间消耗在几个求和计算上，该代码段会随数据变化在每次模拟中运行数千次。通过timeit测试发现，bottleneck.nansum的性能远优于numpy.sum，但实际替换后，整体代码运行速度慢了约5倍，bottleneck.nansum的耗时更是超过原numpy.sum的10倍，与测试结果完全相悖。附上测试性能数据及Profiler运行结果，求分析原因并给出排查方向。

原numpy实现代码 #

import numpy as np
# 生成示例数据
array_size = 100 # 范围100-5000
weighting = np.random.random(array_size)
var_a = np.random.random(array_size)
var_b = np.random.random((array_size, 3))
var_c = np.random.random((array_size, 3))
var_d = np.random.random(array_size)
var_e = np.random.random(array_size)
# 计算结果
res_1 = var_a * weighting
res_2 = np.sum(res_1)
res_3 = np.sum(var_b * var_a[:, np.newaxis], axis=0)
res_4 = np.sum(var_c * var_a[:, np.newaxis], axis=0)
res_5 = np.sum(res_1 / var_e)
res_6 = np.sum(var_d / res_1) / res_2

numpy sum的timeit基准数据 #

36.3 µs ± 579 ns per loop (7次运行，每次10000循环) # array_size=100
103 µs ± 2.88 µs per loop (7次运行，每次10000循环) # array_size=1000
360 µs ± 3.23 µs per loop (7次运行，每次1000循环) # array_size=5000

bottleneck修改后代码 #

res_1 = var_a * weighting
res_2 = bn.nansum(res_1)
res_3 = bn.nansum(var_b * var_a[:, np.newaxis], axis=0)
res_4 = bn.nansum(var_c * var_a[:, np.newaxis], axis=0)
res_5 = bn.nansum(res_1 / var_e)
res_6 = bn.nansum(var_d / res_1) / res_2

bottleneck.nansum的timeit测试数据 #

10.3 µs ± 222 ns per loop (7次运行，每次100000循环) # array_size=100
43 µs ± 1.58 µs per loop (7次运行，每次10000循环) # array_size=1000
177 µs ± 1.66 µs per loop (7次运行，每次10000循环) # array_size=5000

Profiler运行结果 #

ncalls | tottime | percall | cumtime | percall | filename:lineno(function)
83349   525.4   0.006303    525.4   0.006303    ~:0(<built-in method bottleneck.reduce.nansum>)

83349   0.2606  3.127e-06   55.69   0.0006681   fromnumeric.py:2188(sum)
101880  55.43   0.0005441   55.43   0.0005441   ~:0(<method 'reduce' of 'numpy.ufunc' objects>)

分析与排查方向 #

• 数据类型与内存对齐：检查实际运行时数组的数据类型、内存是否对齐。bottleneck对特定dtype（如float32）或未对齐内存的数组可能存在性能损耗，而numpy.sum适配性更强。
• 函数调用开销：timeit测试是批量循环调用，而实际代码是高频次单次调用。bottleneck.nansum的单次调用开销远高于numpy.sum（Profiler显示单次耗时是numpy的近10倍），高频调用下这个差距会被放大。
• bottleneck版本与编译优化：确认bottleneck是否为编译优化版本（如启用MKL/OpenBLAS），部分预编译包可能未开启最优加速，或与当前numpy的BLAS后端不兼容。
• axis参数与维度适配：检查实际求和的axis参数是否与测试一致。bottleneck.nansum在处理特定axis时的优化逻辑可能与numpy不同，比如二维数组axis=0求和是否存在额外内存拷贝。
• 临时数组的影响：实际代码中var_b * var_a[:, np.newaxis]这类临时数组的创建开销是否被误统计到bottleneck耗时中？timeit测试可能复用了数组，而实际代码每次都重新创建。
• 并行策略差异：numpy可能默认启用多线程并行求和，而bottleneck.nansum的并行配置可能不同。检查numpy.show_config()和bottleneck的并行设置，确认并行度差异。

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