numpy ufuncs 速度与循环速度

numpy ufuncs speed vs for loop speed

我读了很多书"avoid for loops with 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 循环:

 In[2]: for i, ti in enumerate(tim):
           values[i] = np.sum(np.sin(prec * ti))

然后,我摆脱了显式 for 循环,并实现了这个:

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

而且这个解决方案对于小型阵列来说速度更快,但是当我按比例放大时,我得到了这样的时间依赖性:

我遗漏了什么或者这是正常行为吗?如果不是,去哪里挖掘?

编辑:根据评论,这里有一些补充信息。时间是用 IPython 的 %timeit%%timeit 测量的,每个 运行 都是在新鲜内核上执行的。我的笔记本电脑是 acer aspire v7-482pg(i7、8GB)。我正在使用:

这是正常的预期行为。将 "avoid for loops with numpy" 语句应用到任何地方都太简单了。如果你正在处理内部循环,它(几乎)总是正确的。但是在外循环的情况下(就像你的情况一样)有更多的例外。特别是如果替代方案是使用广播,因为这通过使用更多内存来加快您的操作。

只是为 "avoid for loops with numpy" 语句添加一些背景知识:

NumPy 数组存储为 类型的连续数组。 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:

%timeit np.sum(arr)
# 10000 loops, best of 3: 24.2 µs per loop

即使您将循环推入 Python C 代码,您离 numpy 性能还很远:

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

这条规则可能会有例外,但这种情况真的很少见。至少只要有一些等效的 numpy 功能。因此,如果您要遍历单个元素,那么您应该使用 numpy。


有时一个简单的 python 循环就足够了。它没有被广泛宣传,但与 Python 函数相比,numpy 函数具有巨大的开销。例如,考虑一个 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,元素很少,因子 2,大约 500 个元素)。但这并不是真正的循环开销,而是函数调用开销!

你的循环解决方案

使用 line-profilerresolution = 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.multiplynp.sinnp.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 调用中花费 500us 时,您不再关心 20us 的开销。

可能需要注意一点: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

再次使用 line_profiler 和 resolution=100:

%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每个步骤花费的时间几乎相同(有些慢一点,有些快一些但总体上非常相似)

另一个影响是当你做乘法时实际的广播变得很重要。即使使用非常智能的 numpy 解决方案,这仍然包括一些额外的计算。对于 resolution=10000,您会看到广播乘法开始占用与循环解决方案相关的更多“% 时间”:

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 数组,请确保您不会 运行 内存不足。

但是关于优化最重要的一点还是:

  • 只有当代码太慢时才优化它!如果速度太慢,则仅在分析代码后才进行优化。

  • 不要盲目相信简化的语句并且再次永远不要在没有分析的情况下进行优化。


最后一个想法:

这些需要循环或广播的功能可以使用 , or if there is no already existing solution in or .

轻松实现

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

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)')