为何NumPy实现比for循环版本更慢？如何优化？

我编写了两个函数np_prod和slice_prod，对形状相同的NumPy数组mat1和mat2执行相同的元素级操作。但测试发现np_prod比slice_prod更慢，测试代码如下：

import numpy as np
import timeit

shape = (10, 600, 600)
mat1 = np.random.rand(*shape).astype(np.float32)
mat2 = np.random.rand(*shape).astype(np.float32)

def np_prod(mat1, mat2):
    res = mat1.copy()
    res *= -1
    res += 1
    res *= (1 - mat2)
    res *= -1
    res += 1
    return res

def slice_prod(mat1, mat2):
    res = mat1.copy()
    for i in range(res.shape[0]):
        res[i] *= -1
        res[i] += 1
        res[i] *= (1 - mat2[i])
        res[i] *= -1
        res[i] += 1
    return res

N = 400
print("For loop product: ", timeit.timeit("slice_prod(mat1, mat2)", setup="from __main__ import slice_prod, mat1, mat2", number=N))
print("Numpy product: ", timeit.timeit("np_prod(mat1, mat2)", setup="from __main__ import np_prod, mat1, mat2", number=N))

测试结果：

For loop product:  15.605191999999999
Numpy product:  18.899588700000002

我的NumPy版本是1.24.2，原本预期NumPy的优化数组操作会更快，但实际for循环版本更优。请问该现象的原因是什么？如何优化np_prod使其更快？（注：示例为简化场景，实际代码中性能差异可达2倍，用time.perf_counter测量）

原因分析 #

• 内存访问与缓存命中率差异：slice_prod每次处理的是(600,600)的连续子数组（NumPy默认采用C顺序存储，子数组在内存中是连续块），缓存命中率更高；而np_prod操作整个大数组时，多次全局读写会占用更多缓存空间，加上临时数组的创建，容易触发缓存失效，增加内存带宽的负载。
• 临时数组的额外开销：np_prod中的(1 - mat2)会创建一个完整的临时数组，占用与原数组相同的内存；而slice_prod中(1 - mat2[i])生成的临时数组仅为子数组大小，内存占用低，垃圾回收和内存分配的开销更小。
• 循环的隐性优化：现代CPU的缓存机制、以及NumPy对切片操作的底层优化，使得对小维度（这里是第一维度的10次循环）的Python循环开销被大幅降低，而多次全局数组操作的累积开销反而超过了循环本身的开销。

优化方案 #

1. 合并运算逻辑，减少内存读写与临时数组 #

原np_prod的运算步骤可以数学合并为单一表达式，避免多次修改数组和创建临时数组：

def optimized_np_prod(mat1, mat2):
    # 直接推导最终表达式：1 - (1 - mat1) * (1 - mat2)
    return 1 - (1 - mat1) * (1 - mat2)

这种方式只需要一次完整的数组运算，大幅减少内存读写次数和临时数组的创建，性能会显著提升。

2. 原地操作时使用out参数避免临时数组 #

如果需要保留原地修改的逻辑，可以使用NumPy的out参数指定输出数组，避免创建中间临时数组：

def optimized_inplace_np_prod(mat1, mat2):
    res = mat1.copy()
    # 1 - mat1 原地计算
    np.subtract(1, res, out=res)
    # 临时存储1 - mat2（或者也可以原地修改mat2，若允许的话）
    temp = 1 - mat2
    np.multiply(res, temp, out=res)
    # 1 - res 原地计算
    np.subtract(1, res, out=res)
    return res

通过out参数复用已有数组内存，减少内存分配和拷贝的开销。

3. 使用Numba进行JIT编译 #

如果合并运算后仍达不到预期性能，可以用Numba对函数进行即时编译，无论向量化还是循环版本都能获得大幅加速：

from numba import jit

@jit(nopython=True)
def numba_prod(mat1, mat2):
    res = mat1.copy()
    for i in range(res.shape[0]):
        res[i] *= -1
        res[i] += 1
        res[i] *= (1 - mat2[i])
        res[i] *= -1
        res[i] += 1
    return res

Numba会将Python代码编译为机器码，消除Python解释器的开销，同时优化内存访问模式，性能通常会远超原生NumPy或纯Python循环。

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