numpy ufuncs速度vs for loop speed

Question

我读过很多“避免与numpy循环”。所以，我试过了。我用的是这个代码（简化版）。一些辅助数据：

 In[1]: import numpy as np
        resolution = 1000                             # this parameter varies
        tim = np.linspace(-np.pi, np.pi, resolution) 
        prec = np.arange(1, resolution + 1)
        prec = 2 * prec - 1
        values = np.zeros_like(tim)

我的第一个实现是使用for循环：

^{pr2}$

然后，我摆脱了显式的for循环，实现了这一点：

 In[3]: values = np.sum(np.sin(tim[:, np.newaxis] * prec), axis=1)

这种解决方案对于小型阵列更快，但当我扩大规模时，我得到了时间依赖性：

我错过了什么还是正常行为？如果不是，去哪里挖？在

编辑：根据评论，以下是一些附加信息。用IPython的%timeit和{}测量时间，每次运行都在新鲜内核上进行。我的笔记本电脑是acer aspire v7-482pg（i7，8GB）。我用的是：

• python 3.5.2版
• numpy 1.11.2+mkl
• Windows 10

Answer 1

这是正常和预期的行为。在任何地方都应用“avoid for loops with numpy”语句太简单了。如果你处理的是内部循环（几乎）总是正确的。但是在外循环的情况下（就像你的例子一样），有更多的例外。尤其是如果另一种选择是使用广播，因为这会通过使用更多的内存来加快操作速度。在

为了给避免for loops with numpy“语句添加一点背景：

NumPy数组存储为具有c类型的连续数组。Pythonint与Cint不同！因此，每当迭代数组中的每个项时，都需要将该项从数组中插入，将其转换为Python int，然后对其执行任何操作，最后可能需要再次将其转换为c整数（称为对值进行装箱和取消装箱）。例如，您希望使用Python sum数组中的项：

import numpy as np
arr = np.arange(1000)
%%timeit
acc = 0
for item in arr:
    acc += item
# 1000 loops, best of 3: 478 µs per loop

你最好用numpy:

^{pr2}$

即使将循环推到Python C代码中，也离numpy性能相差甚远：

%timeit sum(arr)
# 1000 loops, best of 3: 387 µs per loop

这条规则可能会有例外，但这些例外会非常稀少。至少只要有一些等效的numpy功能。因此，如果您要迭代单个元素，那么应该使用numpy。在

---

有时候一个简单的python循环就足够了。虽然没有广为人知，但是numpy函数与Python函数相比有着巨大的开销。例如，考虑一个3元素数组：

arr = np.arange(3)
%timeit np.sum(arr)
%timeit sum(arr)

哪个更快？在

解决方案：Python函数的性能优于numpy解决方案：

# 10000 loops, best of 3: 21.9 µs per loop  <- numpy
# 100000 loops, best of 3: 6.27 µs per loop <- python

---

但这和你的例子有什么关系？事实上并不是很多，因为您总是在数组上使用numpy函数（不是单个元素，甚至不是很少的元素），所以您的内部循环已经使用了优化的函数。这就是为什么两者的性能大致相同（+/-只有很少元素的因子10到大约500个元素的因子2）。但这不是真正的循环开销，而是函数调用开销！在

你的循环解决方案

使用line-profiler和resolution = 100：

def fun_func(tim, prec, values):
    for i, ti in enumerate(tim):
        values[i] = np.sum(np.sin(prec * ti))
%lprun -f fun_func fun_func(tim, prec, values)
Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     1                                           def fun_func(tim, prec, values):
     2       101          752      7.4      5.7      for i, ti in enumerate(tim):
     3       100        12449    124.5     94.3          values[i] = np.sum(np.sin(prec * ti))

95%都花在了循环内部，我甚至把循环体分成几个部分来验证这一点：

def fun_func(tim, prec, values):
    for i, ti in enumerate(tim):
        x = prec * ti
        x = np.sin(x)
        x = np.sum(x)
        values[i] = x
%lprun -f fun_func fun_func(tim, prec, values)
Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     1                                           def fun_func(tim, prec, values):
     2       101          609      6.0      3.5      for i, ti in enumerate(tim):
     3       100         4521     45.2     26.3          x = prec * ti
     4       100         4646     46.5     27.0          x = np.sin(x)
     5       100         6731     67.3     39.1          x = np.sum(x)
     6       100          714      7.1      4.1          values[i] = x

这里的时间消耗者是np.multiply，np.sin，np.sum，通过比较他们的每次通话时间和他们的开销，可以很容易地进行检查：

arr = np.ones(1, float)
%timeit np.sum(arr)
# 10000 loops, best of 3: 22.6 µs per loop

因此，一旦可计算函数调用开销与计算运行时相比很小，您将拥有类似的运行时。即使有100件物品，你的开销也很接近。诀窍是知道他们在哪一点上收支平衡。对于1000个项目，呼叫开销仍然很大：

%lprun -f fun_func fun_func(tim, prec, values)
Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     1                                           def fun_func(tim, prec, values):
     2      1001         5864      5.9      2.4      for i, ti in enumerate(tim):
     3      1000        42817     42.8     17.2          x = prec * ti
     4      1000       119327    119.3     48.0          x = np.sin(x)
     5      1000        73313     73.3     29.5          x = np.sum(x)
     6      1000         7287      7.3      2.9          values[i] = x

但是使用resolution = 5000时，与运行时相比，开销相当低：

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     1                                           def fun_func(tim, prec, values):
     2      5001        29412      5.9      0.9      for i, ti in enumerate(tim):
     3      5000       388827     77.8     11.6          x = prec * ti
     4      5000      2442460    488.5     73.2          x = np.sin(x)
     5      5000       441337     88.3     13.2          x = np.sum(x)
     6      5000        36187      7.2      1.1          values[i] = x

当你在每个np.sin电话中花费500美元时，你不再关心20美元的开销了。在

需要注意的是：line_profiler每行可能包含一些额外的开销，也可能是每个函数调用的开销，因此函数调用开销变为可忽略的点可能会更低！！！在

你的广播解决方案

我从分析第一个解决方案开始，让我们对第二个解决方案进行相同的分析：

def fun_func(tim, prec, values):
    x = tim[:, np.newaxis]
    x = x * prec
    x = np.sin(x)
    x = np.sum(x, axis=1)
    return x

再次使用带有resolution=100的line_profiler：

%lprun -f fun_func fun_func(tim, prec, values)
Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     1                                           def fun_func(tim, prec, values):
     2         1           27     27.0      0.5      x = tim[:, np.newaxis]
     3         1          638    638.0     12.9      x = x * prec
     4         1         3963   3963.0     79.9      x = np.sin(x)
     5         1          326    326.0      6.6      x = np.sum(x, axis=1)
     6         1            4      4.0      0.1      return x

这已经大大超过了开销时间，因此我们比循环快了10倍。在

我还对resolution=1000进行了分析：

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     1                                           def fun_func(tim, prec, values):
     2         1           28     28.0      0.0      x = tim[:, np.newaxis]
     3         1        17716  17716.0     14.6      x = x * prec
     4         1        91174  91174.0     75.3      x = np.sin(x)
     5         1        12140  12140.0     10.0      x = np.sum(x, axis=1)
     6         1           10     10.0      0.0      return x

使用precision=5000：

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     1                                           def fun_func(tim, prec, values):
     2         1           34     34.0      0.0      x = tim[:, np.newaxis]
     3         1       333685 333685.0     11.1      x = x * prec
     4         1      2391812 2391812.0    79.6      x = np.sin(x)
     5         1       280832 280832.0      9.3      x = np.sum(x, axis=1)
     6         1           14     14.0      0.0      return x

1000大小仍然更快，但是正如我们看到的，在循环解决方案中，呼叫开销仍然是不可忽视的。但是对于resolution = 5000来说，每一步所花费的时间几乎是相同的（有些慢一些，有些更快，但总体上非常相似）

另一个影响是>广播当你做乘法运算时，它变得有意义。即使有非常聪明的纽比解决方案，这仍然包括一些额外的计算。对于resolution=10000您可以看到广播乘法相对于循环解决方案开始占用更多的“%time”：

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     1                                           def broadcast_solution(tim, prec, values):
     2         1           37     37.0      0.0      x = tim[:, np.newaxis]
     3         1      1783345 1783345.0    13.9      x = x * prec
     4         1      9879333 9879333.0    77.1      x = np.sin(x)
     5         1      1153789 1153789.0     9.0      x = np.sum(x, axis=1)
     6         1           11     11.0      0.0      return x


Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     8                                           def loop_solution(tim, prec, values):
     9     10001        62502      6.2      0.5      for i, ti in enumerate(tim):
    10     10000      1287698    128.8     10.5          x = prec * ti
    11     10000      9758633    975.9     79.7          x = np.sin(x)
    12     10000      1058995    105.9      8.6          x = np.sum(x)
    13     10000        75760      7.6      0.6          values[i] = x

---

但除了实际花费的时间之外，还有一件事：内存消耗。循环解决方案需要O(n)内存，因为您总是处理n元素。然而，广播解决方案需要O(n*n)内存。如果在循环中使用resolution=20000，则可能需要等待一段时间，但它仍然只需要8bytes/element * 20000 element ~= 160kB，但对于广播，则需要~3GB。而这忽略了常数因子（比如临时数组又称中间数组）！如果你再往前走，你的记忆会很快用完的！在

---

再次总结要点：

• 如果你对numpy数组中的单个项执行python循环，那么这是错误的。在
• 如果在numpy数组的子数组上循环，请确保每个循环中的函数调用开销与在函数中花费的时间相比是可以忽略的。在
• 如果你广播numpy数组，确保你不会耗尽内存。在

但优化最重要的一点仍然是：

• 只有当代码太慢时才进行优化！如果速度太慢，那么只有在分析代码之后才进行优化。

• 不要盲目地相信简化的语句，同样，永远不要在没有分析的情况下进行优化。

---

最后一个想法是：

如果numpy或scipy中没有现有的解决方案，则可以使用cython、numba或numexpr轻松实现此类需要循环或广播的函数。在

例如，numba函数将循环解决方案的内存效率与低resolutions的广播解决方案的速度结合起来，如下所示：

from numba import njit

import math

@njit
def numba_solution(tim, prec, values):
    size = tim.size
    for i in range(size):
        ti = tim[i]
        x = 0
        for j in range(size):
            x += math.sin(prec[j] * ti)
        values[i] = x

正如评论中所指出的，numexpr还可以非常快速地评估广播计算，，而不需要内存O(n*n)内存：

>>> import numexpr
>>> tim_2d = tim[:, np.newaxis]
>>> numexpr.evaluate('sum(sin(tim_2d * prec), axis=1)')