For 循环比纯 numpy 函数更快

Question

[披露：在写这个问题的过程中，我发现了发生了什么，所以我将提出我自己的解决方案。如果您有任何补充，我很乐意阅读其他一些建议。]

我正在研究一种机器学习方法，该方法涉及在小图像裁剪上多次运行数学和图像处理函数。目前我想要实现的是通过优化这些函数来使管道更快。最初，每个训练样本都是在 for 循环中单独处理的，所以我认为一个简单的改进是将 2D numpy 数组（图像）组合成一个 3D 数组，然后一起处理它们，而不需要 for 循环。我进行了一些性能测量，我惊讶地发现 for 循环中的处理实际上比纯 numpy 计算更快。

这就是我正在做的（请在最后找到完整的脚本）。我用 timeit 进行了测量。设置代码如下（

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
  

• 对整个数组进行操作（

  "pure np"

  ）：

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

• 无需复制(

  "pure np (no copy)"

  ):

  out = (a+b)/2
  

我对 sqrt 做同样的事情；为了完整起见：

• 循环：

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

• numpy：

  out[:] = np.sqrt(a)
  

• 无副本：

  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

、

multipy

、

log1p

和一些 cv2 函数）进行了相同的测试，结果都相同。我无法理解这一点，因为最基本的 Python 编程原则是使用 Numpy，因为它比循环快得多。有谁知道为什么会发生这种情况？更重要的是：如果省略循环没有帮助，我可以做什么来加速我的脚本？

这是我用于测试的完整脚本（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")

Answer 1

分配“大”数组似乎是罪魁祸首。我可以通过使用 numpy 函数的

out

参数或就地进行计算来加速这些计算：

np.add(a, b, out=out)
out /= 2

表示平均值，并且

np.sqrt(a, out=out)

对于开方。

我可以对循环版本做类似的事情来提高速度：

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

和

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

新结果：

=========mean==========
    np loop:  3193.34 ms
    pure np:  8400.82 ms
    pure np (no copy):  7102.62 ms
    inplace np:  1165.70 ms
    np inplace loop:  2167.79 ms
=========sqrt==========
    np loop:  1428.44 ms
    pure np:  4185.10 ms
    pure np (no copy):  3459.68 ms
    inplace np:  588.06 ms
    np inplace loop:  860.07 ms

对于

arr_size=(100,100,100)

：

=========mean==========
    np loop:  11747.72 ms
    pure np:  27096.73 ms
    pure np (no copy):  23890.63 ms
    inplace np:  5679.39 ms
    np inplace loop:  8931.12 ms
=========sqrt==========
    np loop:  5062.36 ms
    pure np:  13357.57 ms
    pure np (no copy):  10620.47 ms
    inplace np:  2488.66 ms
    np inplace loop:  3279.80 ms

对于 `arr_size=(30, 300, 300)：

=========mean==========
    np loop:  24232.15 ms
    pure np:  78151.40 ms
    pure np (no copy):  64296.45 ms
    inplace np:  21251.48 ms
    np inplace loop:  23041.19 ms
=========sqrt==========
    np loop:  14685.71 ms
    pure np:  37801.05 ms
    pure np (no copy):  30844.35 ms
    inplace np:  9997.73 ms
    np inplace loop:  10121.31 ms