为何For循环比纯Numpy运算更快？求性能优化方案

说明：撰写此问题时我已了解背后原因，将提出自己的解决方案，欢迎各位补充建议。

我正在开发一种机器学习方法，需要对小图像块反复执行大量数学及图像处理函数。为加速流程，我原本计划将逐个处理的2D Numpy数组（图像）合并为单个3D数组，去除For循环用纯Numpy批量处理，但性能测试结果却让我震惊：For循环处理反而比纯Numpy运算更快。

初始化测试代码（init_script） #

import numpy as np
import cv2
a = np.random.rand(30, 100, 100).astype(np.float32)
b = np.random.rand(30, 100, 100).astype(np.float32)
out = np.empty_like(a)

测试的三种实现方式 #

均值运算 #

• 循环方式（np loop）：

for i, (a_slice, b_slice) in enumerate(zip(a, b)):
    out[i] = (a_slice+b_slice)/2

• 纯Numpy批量处理（pure np）：

out[:] = (a+b)/2

• 纯Numpy无拷贝版本（pure np (no copy)）：

out = (a+b)/2

平方根运算 #

• 循环方式：

for i, a_slice in enumerate(a):
    out[i] = np.sqrt(a_slice)

• 纯Numpy批量处理：

out[:] = np.sqrt(a)

• 纯Numpy无拷贝版本：

out = np.sqrt(a)

测试结果 #

10000次运行基础测试 #

=========mean==========
    np loop:  2939.34 ms
    pure np:  8441.93 ms
    pure np (no copy):  7417.07 ms
=========sqrt==========
    np loop:  1353.00 ms
    pure np:  4304.14 ms
    pure np (no copy):  3546.86 ms

增加样本量（arr_size = (100, 100, 100)） #

=========mean==========
    np loop:  11125.34 ms
    pure np:  26596.88 ms
    pure np (no copy):  24165.19 ms
=========sqrt==========
    np loop:  5107.81 ms
    pure np:  13542.24 ms
    pure np (no copy):  10445.66 ms

增大图像尺寸（arr_size = (30, 300, 300)） #

=========mean==========
    np loop:  24655.34 ms
    pure np:  72427.00 ms
    pure np (no copy):  66605.25 ms
=========sqrt==========
    np loop:  16771.90 ms
    pure np:  38191.14 ms
    pure np (no copy):  30087.40 ms

我还测试了maximum、multiply、log1p及部分cv2函数，结果一致。这与“使用Numpy替代循环以大幅提速”的Python基本编程原则相悖，请问这是什么原因？若去除循环无法提速，我该如何优化脚本？

完整测试脚本（Numpy版本：1.26.4） #

import timeit

arr_size = (30, 100, 100)  # 30 images of size 100x100

init_script = f"import numpy as np;" \
              f"a = np.random.rand(*{arr_size}).astype(np.float32);" \
              f"b = np.random.rand(*{arr_size}).astype(np.float32);" \
              f"out = np.empty_like(a)"

mean_tests = {
    "np loop": "for i, (a_slice, b_slice) in enumerate(zip(a, b)): out[i] = (a_slice+b_slice)/2",
    "pure np": "out[:] = (a+b)/2",
    "pure np (no copy)": "out = (a+b)/2",
}

sqrt_tests = {
    "np loop": "for i, a_slice in enumerate(a): out[i] = np.sqrt(a_slice)",
    "pure np": "out[:] = np.sqrt(a)",
    "pure np (no copy)": "out = np.sqrt(a)",
}

tests = {"mean": mean_tests, "sqrt": sqrt_tests}

for func, test in tests.items():
    print(f"========={func}==========")
    for test_case, cmd in test.items():
        elapsed = timeit.timeit(cmd, init_script, number=10000)
        print(f"\t{test_case}: {elapsed*1000: .2f} ms")

原因分析 #

这种反直觉结果的核心原因是：

1. 中间数组开销：Numpy批量运算会生成大尺寸临时数组（比如(a+b)/2会先创建完整的3D数组a+b，再生成除以2的结果数组），两次大内存分配和拷贝的开销远超过Python循环的解释器开销。
2. CPU缓存命中率：循环处理小2D切片时，数据更容易被CPU缓存容纳，运算效率更高；而3D批量处理需要操作更大的连续内存块，内存带宽压力更大。
3. 小数组优化不足：Numpy底层对小尺寸数组的批量处理没有针对性优化，反而因为统一的内存管理逻辑带来额外开销。

优化方案 #

1. 用Numba加速循环 #

使用numba.jit将Python编译为机器码，消除解释器开销，同时保留小切片的缓存优势：

from numba import jit

@jit(nopython=True)
def mean_loop_numba(a, b, out):
    for i in range(a.shape[0]):
        out[i] = (a[i] + b[i]) / 2

@jit(nopython=True)
def sqrt_loop_numba(a, out):
    for i in range(a.shape[0]):
        out[i] = np.sqrt(a[i])

2. 减少Numpy中间数组 #

使用Numpy的原地运算方法，避免创建大临时数组：

# 均值运算：先原地相加，再原地除以2，无临时数组
np.add(a, b, out=out)
np.divide(out, 2, out=out)

# 平方根运算：原地修改输出数组
np.sqrt(a, out=out)

3. 确保数组内存连续 #

非连续数组会降低Numpy运算效率，可通过以下方式确保内存布局为C连续：

a = np.ascontiguousarray(a, dtype=np.float32)
b = np.ascontiguousarray(b, dtype=np.float32)

4. 替换为OpenCV批量函数 #

OpenCV的图像处理函数对3D数组的底层优化更高效，比如直接用cv2.sqrt替代np.sqrt：

cv2.sqrt(a, out=out)

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