python - 如何矢量化循环遍历 3D 点数组的 python 函数？

Question

代码如下:

import numpy as np
from numpy.random import random

@profile
def point_func(point, points, funct):
    return np.sum(funct(np.sqrt(((point - points)**2)).sum(1)))

@profile
def point_afunc(ipoints, epoints, funct):
    res = np.zeros(len(ipoints))
    for idx, point in enumerate(ipoints):
        res[idx] = point_func(point, epoints, funct)
    return res

@profile
def main():
    points = random((5000,3))
    rpoint = random((1,3))

    pres = point_func(rpoint, points, lambda r : r**3)

    ares = point_afunc(points, points, lambda r : r**3)

if __name__=="__main__":
    main()

我已经用 kernprof 分析了它并得到了这个:

Timer unit: 1e-06 s

Total time: 2.25667 s File: point-array-vectorization.py Function: point_func at line 4

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     4                                           @profile
     5                                           def point_func(point, points, funct):
     6      5001    2256667.0    451.2    100.0      return np.sum(funct(np.sqrt(((point - points)**2)).sum(1)))

Total time: 2.27844 s File: point-array-vectorization.py Function: point_afunc at line 8

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     8                                           @profile
     9                                           def point_afunc(ipoints, epoints, funct):
    10         1          5.0      5.0      0.0      res = np.zeros(len(ipoints))
    11      5001       4650.0      0.9      0.2      for idx, point in enumerate(ipoints):
    12      5000    2273789.0    454.8     99.8          res[idx] = point_func(point, epoints, funct)
    13         1          0.0      0.0      0.0      return res

Total time: 2.28239 s File: point-array-vectorization.py Function: main at line 15

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    15                                           @profile
    16                                           def main():
    17         1        145.0    145.0      0.0      points = random((5000,3))
    18         1          2.0      2.0      0.0      rpoint = random((1,3))
    19                                           
    20         1        507.0    507.0      0.0      pres = point_func(rpoint, points, lambda r : r**3)
    21                                           
    22         1    2281731.0 2281731.0    100.0      ares = point_afunc(points, points, lambda r : r**3)

所以这部分花费了大部分时间:

11      5001       4650.0      0.9      0.2      for idx, point in enumerate(ipoints):
    12      5000    2273789.0    454.8     99.8          res[idx] = point_func(point, epoints, funct)

我想看看时间损失是不是因为在for循环中调用了funct造成的。为此，我想使用 numpy.vectorize 向量化 point_afunc。我试过了，但它似乎把点矢量化了:循环最终循环遍历各个点组件。

@profile
def point_afunc(ipoints, epoints, funct):
    res = np.zeros(len(ipoints))
    for idx, point in enumerate(ipoints):
        res[idx] = point_func(point, epoints, funct)
    return res

point_afunc = np.vectorize(point_afunc)

导致错误:

  File "point-array-vectorization.py", line 24, in main
    ares = point_afunc(points, points, lambda r : r**3)
  File "/usr/lib/python3.6/site-packages/numpy/lib/function_base.py", line 2755, in __call__
    return self._vectorize_call(func=func, args=vargs)
  File "/usr/lib/python3.6/site-packages/numpy/lib/function_base.py", line 2825, in _vectorize_call
    ufunc, otypes = self._get_ufunc_and_otypes(func=func, args=args)
  File "/usr/lib/python3.6/site-packages/numpy/lib/function_base.py", line 2785, in _get_ufunc_and_otypes
    outputs = func(*inputs)
  File "/usr/lib/python3.6/site-packages/line_profiler.py", line 115, in wrapper
    result = func(*args, **kwds)
  File "point-array-vectorization.py", line 10, in point_afunc
    res = np.zeros(len(ipoints))
TypeError: object of type 'numpy.float64' has no len()

不知何故，不是对ipoints 中的每个点 进行矢量化，而是对点的分量进行矢量化？

编辑:尝试了下面@John Zwinck 的建议并使用了 numba。我使用 @jit 的执行时间比没有它时更长。如果我从所有函数中删除 @profile 装饰器，并将其替换为 @jit 用于 point_func 和 point_afunc，这些是执行时间:

time ./point_array_vectorization.py 

real    0m3.686s
user    0m3.584s
sys 0m0.077s
point-array-vectorization> time ./point_array_vectorization.py 

real    0m3.683s
user    0m3.596s
sys 0m0.063s
point-array-vectorization> time ./point_array_vectorization.py 

real    0m3.751s
user    0m3.658s
sys 0m0.070s

并删除所有 @jit 装饰器:

point-array-vectorization> time ./point_array_vectorization.py 

real    0m2.925s
user    0m2.874s
sys 0m0.030s
point-array-vectorization> time ./point_array_vectorization.py 

real    0m2.950s
user    0m2.902s
sys 0m0.029s
point-array-vectorization> time ./point_array_vectorization.py 

real    0m2.951s
user    0m2.886s
sys 0m0.042s

我是否需要更多帮助 numba 编译器？

编辑:point_afunc 是否可以使用 numpy 以某种方式在没有 for 循环的情况下编写？

编辑:将循环版本与Peter的numpy广播版本进行比较，循环版本更快:

Timer unit: 1e-06 s

Total time: 2.13361 s
File: point_array_vectorization.py
Function: point_func at line 7

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     7                                           @profile
     8                                           def point_func(point, points, funct):
     9      5001    2133615.0    426.6    100.0      return np.sum(funct(np.sqrt(((point - points)**2)).sum(1)))

Total time: 2.1528 s
File: point_array_vectorization.py
Function: point_afunc at line 11

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    11                                           @profile
    12                                           def point_afunc(ipoints, epoints, funct):
    13         1          5.0      5.0      0.0      res = np.zeros(len(ipoints))
    14      5001       4176.0      0.8      0.2      for idx, point in enumerate(ipoints):
    15      5000    2148617.0    429.7     99.8          res[idx] = point_func(point, epoints, funct)
    16         1          0.0      0.0      0.0      return res

Total time: 2.75093 s
File: point_array_vectorization.py
Function: new_point_afunc at line 18

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    18                                           @profile
    19                                           def new_point_afunc(ipoints, epoints, funct):
    20         1    2750926.0 2750926.0    100.0      return np.sum(funct(np.sqrt((ipoints[:, None, :] - epoints[None, :, :])**2).sum(axis=-1)), axis=1)

Total time: 4.90756 s
File: point_array_vectorization.py
Function: main at line 22

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    22                                           @profile
    23                                           def main():
    24         1        170.0    170.0      0.0      points = random((5000,3))
    25         1          4.0      4.0      0.0      rpoint = random((1,3))
    26         1        546.0    546.0      0.0      pres = point_func(rpoint, points, lambda r : r**3)
    27         1    2155829.0 2155829.0     43.9      ares = point_afunc(points, points, lambda r : r**3)
    28         1    2750945.0 2750945.0     56.1      vares = new_point_afunc(points, points, lambda r : r**3)
    29         1         71.0     71.0      0.0      assert(np.max(np.abs(ares-vares)) < 1e-15)

Answer 1

numpy.vectorize() 在性能方面没有任何用处:它只是构建隐藏的 Python for 循环的语法糖(或者更确切地说，语法氰化物)。它不会帮助你。

可能对您有很大帮助的一件事是 Numba .它可以即时编译您的原始代码，并且可能会大大加快速度。只需将 @profile 装饰器替换为 @numba.jit。