为什么`frompyfunc`比`vectorize`表现更好？

Question

Numpy提供了vectorize和frompyfunc，它们具有类似的功能。

正如这个SO-post所指出的那样，vectorize封装了frompyfunc并正确处理返回数组的类型，而frompyfunc返回一个np.object类型的数组。

然而，frompyfunc在所有大小的情况下始终比vectorize快10-20％，这也不能用不同的返回类型来解释。

考虑以下变体：

import numpy as np

def do_double(x):
    return 2.0*x

vectorize = np.vectorize(do_double)

frompyfunc = np.frompyfunc(do_double, 1, 1)

def wrapped_frompyfunc(arr):
    return frompyfunc(arr).astype(np.float64)

"wrapped_frompyfunc只是将frompyfunc的结果转换为正确的类型 - 我们可以看到，这个操作的成本几乎可以忽略不计。它产生了以下时间（蓝线是frompyfunc）："

我会预期`vectorize`有更多的开销——但这只有在小规模时才能看到。另一方面，将`np.object`转换为`np.float64`也是在`wrapped_frompyfunc`中完成的——但仍然比较快。

这种性能差异如何解释？

---

使用perfplot包生成时间比较的代码（给定上述函数）：

import numpy as np
import perfplot
perfplot.show(
    setup=lambda n: np.linspace(0, 1, n),
    n_range=[2**k for k in range(20,27)],
    kernels=[
        frompyfunc, 
        vectorize, 
        wrapped_frompyfunc,
        ],
    labels=["frompyfunc", "vectorize", "wrapped_frompyfunc"],
    logx=True,
    logy=False,
    xlabel='len(x)',
    equality_check = None,  
    )

---

注意：对于更小的尺寸，vectorize 的开销要高得多，但这是可以预料的（毕竟它包装了 frompyfunc）：

- ead

---

注意事项

vectorize函数主要是为了方便而提供的，而不是为了性能。实现本质上是一个for循环。 如果未指定otypes，则将使用带有第一个参数的函数调用来确定输出的数量。如果cache为True，则会缓存此调用的结果，以防止多次调用该函数。然而，为了实现缓存，必须包装原始函数，这将减慢后续调用的速度，因此只有在您的函数很耗时时才这样做。 - Giacomo Alzetta

从你的陈述中“*vectorize包装了frompyfunc 和 …*”，我觉得很明显vectorize更慢。vectorize（1）调用frompyfunc并且另外（2）做一些其他的事情。为什么你期望做（1）+（2）比只做（2）更快呢？ - Socowi

@Socowi 因为我正在使用 wrapped_frompyfunc 进行 (2) 操作，而且它比 vectorize 要快得多。 - ead

你的“封装函数”根本无法与vectorize所做的相提并论。它也无法与vectorize所带来的额外开销相比。当然，这个例子比vectorize更快。 - user3483203

我认为你需要详细检查vectorize代码。仅凭推理和时间不足够。 - hpaulj

Answer 1

根据@hpaulj的提示，我们可以对vectorize函数进行剖析：

arr=np.linspace(0,1,10**7)
%load_ext line_profiler

%lprun -f np.vectorize._vectorize_call \
       -f np.vectorize._get_ufunc_and_otypes  \
       -f np.vectorize.__call__  \
       vectorize(arr)

这表明在_vectorize_call中花费了100%的时间：

Timer unit: 1e-06 s

Total time: 3.53012 s
File: python3.7/site-packages/numpy/lib/function_base.py
Function: __call__ at line 2063

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
  2063                                               def __call__(self, *args, **kwargs):
  ...                                         
  2091         1    3530112.0 3530112.0    100.0          return self._vectorize_call(func=func, args=vargs)

...

Total time: 3.38001 s
File: python3.7/site-packages/numpy/lib/function_base.py
Function: _vectorize_call at line 2154

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
  2154                                               def _vectorize_call(self, func, args):
  ...
  2161         1         85.0     85.0      0.0              ufunc, otypes = self._get_ufunc_and_otypes(func=func, args=args)
  2162                                           
  2163                                                       # Convert args to object arrays first
  2164         1          1.0      1.0      0.0              inputs = [array(a, copy=False, subok=True, dtype=object)
  2165         1     117686.0 117686.0      3.5                        for a in args]
  2166                                           
  2167         1    3089595.0 3089595.0     91.4              outputs = ufunc(*inputs)
  2168                                           
  2169         1          4.0      4.0      0.0              if ufunc.nout == 1:
  2170         1     172631.0 172631.0      5.1                  res = array(outputs, copy=False, subok=True, dtype=otypes[0])
  2171                                                       else:
  2172                                                           res = tuple([array(x, copy=False, subok=True, dtype=t)
  2173                                                                        for x, t in zip(outputs, otypes)])
  2174         1          1.0      1.0      0.0          return res

这表明我在假设中忽略了一个部分：双数组在预处理步骤中完全转换为对象数组（从内存角度来看不是很明智）。对于wrapped_frompyfunc，其他部分类似：

Timer unit: 1e-06 s

Total time: 3.20055 s
File: <ipython-input-113-66680dac59af>
Function: wrapped_frompyfunc at line 16

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    16                                           def wrapped_frompyfunc(arr):
    17         1    3014961.0 3014961.0     94.2      a = frompyfunc(arr)
    18         1     185587.0 185587.0      5.8      b = a.astype(np.float64)
    19         1          1.0      1.0      0.0      return b

当我们查看峰值内存消耗时（例如通过/usr/bin/time python script.py），我们会发现，矢量化（vectorized）版本的内存消耗是使用更复杂策略的frompyfunc的两倍：双数组以NPY_BUFSIZE（即8192）的块处理，并且在同一时间内仅有8192个Python-floats（24字节+8字节指针）存在于内存中（而不是数组中的元素数量，这可能要高得多）。从操作系统中保留内存和更多缓存未命中的成本可能导致运行时间更长。

我的收获：

• 预处理步骤将所有输入转换为对象数组可能根本不需要，因为frompyfunc具有更复杂的处理方式。
• 当结果的ufunc应在“真实代码”中使用时，既不应使用vectorize也不应使用frompyfunc。 相反，应编写C代码或使用numba /类似工具。

---

在对象数组上调用frompyfunc所需的时间少于在双精度数组上调用的时间：

arr=np.linspace(0,1,10**7)
a = arr.astype(np.object)
%timeit frompyfunc(arr)  # 1.08 s ± 65.8 ms
%timeit frompyfunc(a)    # 876 ms ± 5.58 ms

然而，上述的line-profiler-timings并没有显示出使用ufunc比双精度浮点数对象更有优势：3.089595秒与3014961.0秒。我怀疑这是因为当所有对象都被创建时，会有更多的缓存未命中，而只有8192个（256Kb）热对象在L2缓存中。

- ead

Answer 2

问题完全无意义。如果速度是问题，那么向量化和frompyfunc都不是答案。它们之间的任何速度差异都相形见绌，与更快的方法相比毫不重要。

我发现这个问题在想为什么frompyfunc破坏了我的代码（它返回对象），而向量化起作用了（它返回我告诉它要做的事情），并发现人们谈论速度。

现在，在2020年代，numba/jit可用，它可以轻松地打破frompyfunc的任何速度优势。

我编写了一个玩具应用程序，从另一个应用程序返回一个大型np.uint8数组，并获得了以下结果。

pure python                      200 ms
vectorize                         58 ms
frompyfunc + cast back to uint8   53 ms   
np.empty + numba/njit             55 us (4 cores, 100 us single core)

所以速度比numpy快1000倍，比纯python快4000倍

如果有人在意的话，我可以发布代码。编写njit版本只需要在纯python函数之前添加@njit这一行，所以你不需要特别强大的技术来完成。

与将函数包装在vectorize中相比，它不太方便，因为你需要手动编写循环遍历numpy数组的内容，但它可以避免编写外部C函数。你需要用numpy/C类似的python子集进行编写，避免使用python对象。

也许我对numpy有些苛刻，在要求它向量化纯python函数方面。那么，如果我用numba对比benchmark一个像min这样的numpy本地数组函数呢？

令人惊叹的是，在一个385x360的np.uint8数组上，使用numba/jit比np.min快10倍。np.min(array)的基准时间是230微秒。Numba在单核情况下达到了60微秒，并且在四核情况下达到了22微秒。

# minimum graphical reproducible case of difference between
# frompyfunc and vectorize

# apparently, from stack overflow,
# vectorize returns correct type, but is slow
# frompyfunc always returns an object

# let's see which is faster, casting frompyfunc, or plain vectorize
# and then compare those with plain python, and with njit

# spoiler
# python       200 ms
# vectorise     58 ms
# frompyfunc    53 ms
# njit parallel 55 us


from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import sys
import time
from numba import njit, prange

THRESH = 128
FNAME = '3218_M.PNG' # monochrome screen grab of a sudoku puzzle
ROW = 200

def th_python(x, out):
    rows, cols = x.shape
    for row in range(rows):
        for col in range(cols):
            val = 250
            if x[row, col]<THRESH:
                val = 5
            out[row, col] = val

@njit(parallel=True)
def th_jit(x, out):
    rows, cols = x.shape
    for row in prange(rows):
        for col in prange(cols):
            val = 250
            if x[row, col]<THRESH:
                val = 5
            out[row, col] = val

@njit(parallel=True)
def min_jit(x):
    rows, cols = x.shape
    minval = 255
    for row in prange(rows):
        for col in prange(cols):
            val = x[row, col]
            if val<minval:
                minval = val
    return minval
           

def threshold(x):
    out = 250
    if x<THRESH:
        out = 5
    return np.uint8(out)

th_fpf = np.frompyfunc(threshold,1,1)
th_vec = np.vectorize(threshold, otypes=[np.uint8])

# load an image
image = Image.open(FNAME)
print(f'{image.mode=}')
npim = np.array(image)

# see what we've got
print(f'{type(npim)=}')
print(f'{type(npim[0,0])=}')
# print(npim[ROW,:])
print(f'{npim.shape=}')
print(f'{sys.getsizeof(npim)=}')

# plt.imshow(npim, cmap='gray', vmin=0, vmax=255)
# plt.show()

# threshold it with plain python
start = time.time()
npimpp = np.empty(npim.shape, dtype=np.uint8)
alloc = time.time()
th_python(npim, npimpp)
done = time.time()
print(f'\nallocation took {alloc-start:g} seconds')
print(f'computation took {done-alloc:g} seconds')
print(f'total plain python took {done-start:g} seconds')
print(f'{sys.getsizeof(npimpp)=}')

# use vectorize
start = time.time()
npimv = th_vec(npim)
done = time.time()
print(f'\nvectorize took {done-start:g} seconds')
print(f'{sys.getsizeof(npimv)=}')

# use fpf followed by cast
start = time.time()
npimf = th_fpf(npim)
comp = time.time()
npimfc = np.array(npimf, dtype=np.uint8)
done = time.time()
print(f'\nfunction took {comp-start:g} seconds')
print(f'cast took {done-comp:g} seconds')
print(f'total was {done-start:g} seconds')
print(f'{sys.getsizeof(npimf)=}')

# threshold it with numba jit
for i in range(2):
    print(f'\n go number {i}')
    start = time.time()
    npimjit = np.empty(npim.shape, dtype=np.uint8)
    alloc = time.time()
    th_jit(npim, npimjit)
    done = time.time()
    print(f'\nallocation took {alloc-start:g} seconds')
    print(f'computation took {done-alloc:g} seconds')
    print(f'total with jit took {done-start:g} seconds')
    print(f'{sys.getsizeof(npimjit)=}')

# what about if we use a numpy native function?
start = time.time()
npmin = np.min(npim)
done = time.time()
print(f'\ntotal for np.min was {done-start:g} seconds')

for i in range(2):
    print(f'\n go number {i}')
    start = time.time()
    jit_min = min_jit(npim)
    done = time.time()
    print(f'total with min_jit took {done-start:g} seconds')

- Neil_UK

---

谢谢你分享这个。我对代码很感兴趣。 - jamesbond007

@jamesbond007 给你，我成功找到了。 - Neil_UK