【问题标题】:Vectorize a 6 for loop cumulative sum in python在python中向量化一个6 for循环累积和
【发布时间】:2019-04-17 17:33:27
【问题描述】:

数学题是:

总和中的表达式实际上比上面的要复杂得多,但这是一个最小的工作示例,不会使事情过于复杂。我已经使用 6 个嵌套的 for 循环在 Python 中编写了此代码,并且正如预期的那样,它的性能非常糟糕(真实形式的性能很差,需要评估数百万次),即使在 Numba、Cython 和朋友的帮助下也是如此。这里是使用嵌套的 for 循环和累积和编写的:

import numpy as np

def func1(a,b,c,d):
    '''
    Minimal working example of multiple summation
    '''
    B = 0
    for ai in range(0,a):
        for bi in range(0,b):
            for ci in range(0,c):
                for di in range(0,d):
                    for ei in range(0,ai+bi):
                        for fi in range(0,ci+di):
                            B += (2)**(ei-fi-ai-ci-di+1)*(ei**2-2*(ei*fi)-7*di)*np.math.factorial(ei)


    return a, b, c, d, B

表达式由 4 个数字作为输入来控制,对于 func1(4,6,3,4)B 的输出是 21769947.844726562。

我环顾四周寻求帮助,并找到了几篇 Stack 帖子,其中一些示例如下:

Exterior product in NumPy : Vectorizing six nested loops

Vectorizing triple for loop in Python/Numpy with different array shapes

Python vectorizing nested for loops

我尝试使用从这些有用的帖子中学到的东西,但经过多次尝试,我一直得到错误的答案。即使对一个内部总和进行矢量化,也会为真正的问题带来巨大的性能提升,但总和范围不同的事实似乎让我失望了。有没有人有任何关于如何在这方面取得进展的提示?

【问题讨论】:

  • 您可以通过使用查找表记住阶乘调用来节省一些时间。除非np 已经实现了
  • 这并不能回答您的问题,但如果您关心的是速度,您就不能用 C++ 编写这部分代码吗?我预计至少有 10 倍的加速,但即使是 100 倍也不会让我感到惊讶。
  • 好像没问题。预期的答案是什么?
  • 你试过 PyPy 吗?
  • 缓存 np.math.factorial 在我的机器中提供了 30% 的加速。 cache_fat_ei = [ np.math.factorial(ei) for ei in range(0,a+b) ].

标签: python numpy for-loop vectorization


【解决方案1】:

编辑 3:

最终(我认为)版本,更简洁、更快速地结合了来自 max9111's answer 的想法。

import numpy as np
from numba import as nb

@nb.njit()
def func1_jit(a, b, c, d):
    # Precompute
    exp_min = 5 - (a + b + c + d)
    exp_max = b
    exp = 2. ** np.arange(exp_min, exp_max + 1)
    fact_e = np.empty((a + b - 2))
    fact_e[0] = 1
    for ei in range(1, len(fact_e)):
        fact_e[ei] = ei * fact_e[ei - 1]
    # Loops
    B = 0
    for ai in range(0, a):
        for bi in range(0, b):
            for ci in range(0, c):
                for di in range(0, d):
                    for ei in range(0, ai + bi):
                        for fi in range(0, ci + di):
                            B += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]
    return B

这已经比以前的任何选项都快了,但我们仍然没有利用多个 CPU。一种方法是在函数本身内,例如并行化外循环。这会在每次调用创建线程时增加一些开销,因此对于较小的输入实际上会慢一些,但对于较大的值应该会明显更快:

import numpy as np
from numba import as nb

@nb.njit(parallel=True)
def func1_par(a, b, c, d):
    # Precompute
    exp_min = 5 - (a + b + c + d)
    exp_max = b
    exp = 2. ** np.arange(exp_min, exp_max + 1)
    fact_e = np.empty((a + b - 2))
    fact_e[0] = 1
    for ei in range(1, len(fact_e)):
        fact_e[ei] = ei * fact_e[ei - 1]
    # Loops
    B = np.empty((a,))
    for ai in nb.prange(0, a):
        Bi = 0
        for bi in range(0, b):
            for ci in range(0, c):
                for di in range(0, d):
                    for ei in range(0, ai + bi):
                        for fi in range(0, ci + di):
                            Bi += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]
        B[ai] = Bi
    return np.sum(B)

或者,如果您有很多点要评估函数,您也可以在该级别进行并行化。这里a_arrb_arrc_arrd_arr 是要评估函数的值的向量:

from numba import as nb

@nb.njit(parallel=True)
def func1_arr(a_arr, b_arr, c_arr, d_arr):
    B_arr = np.empty((len(a_arr),))
    for i in nb.prange(len(B_arr)):
        B_arr[i] = func1_jit(a_arr[i], b_arr[i], c_arr[i], d_arr[i])
    return B_arr

最佳配置取决于您的输入、使用模式、硬件等,因此您可以结合不同的想法来适应您的情况。


编辑 2:

其实,忘记我之前说的话吧。最好的办法是对算法进行 JIT 编译,但以更有效的方式。首先计算昂贵的部分(我取了指数和阶乘),然后将其传递给编译后的循环函数:

import numpy as np
from numba import njit

def func1(a, b, c, d):
    exp_min = 5 - (a + b + c + d)
    exp_max = b
    exp = 2. ** np.arange(exp_min, exp_max + 1)
    ee = np.arange(a + b - 2)
    fact_e = scipy.special.factorial(ee)
    return func1_inner(a, b, c, d, exp_min, exp, fact_e)

@njit()
def func1_inner(a, b, c, d, exp_min, exp, fact_e):
    B = 0
    for ai in range(0, a):
        for bi in range(0, b):
            for ci in range(0, c):
                for di in range(0, d):
                    for ei in range(0, ai + bi):
                        for fi in range(0, ci + di):
                            B += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]
    return B

在我的实验中,这是迄今为止最快的选项,并且占用的额外内存很少(仅预先计算的值,大小与输入呈线性关系)。

a, b, c, d = 4, 6, 3, 4
# The original function
%timeit func1_orig(a, b, c, d)
# 2.07 ms ± 33.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# The grid-evaluated function
%timeit func1_grid(a, b, c, d)
# 256 µs ± 25 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
# The precompuation + JIT-compiled function
%timeit func1_jit(a, b, c, d)
# 19.6 µs ± 3.25 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)

嗯,总是可以选择对整个事物进行网格评估:

import numpy as np
import scipy.special

def func1(a, b, c, d):
    ai, bi, ci, di, ei, fi = np.ogrid[:a, :b, :c, :d, :a + b - 2, :c + d - 2]
    # Compute
    B = (2.) ** (ei - fi - ai - ci - di + 1) * (ei ** 2 - 2 * (ei * fi) - 7 * di) * scipy.special.factorial(ei)
    # Mask out of range elements for last two inner loops
    m = (ei < ai + bi) & (fi < ci + di)
    return np.sum(B * m)

print(func1(4, 6, 3, 4))
# 21769947.844726562

我使用了scipy.special.factorial,因为显然np.factorial 出于某种原因不适用于数组。

显然,随着参数的增加,其内存消耗会非常快速增长。代码实际上执行了比必要更多的计算,因为两个内部循环具有不同数量的迭代,因此(在此方法中)您必须使用最大的,然后删除不需要的。希望矢量化能够弥补这一点。一个小的 IPython 基准测试:

a, b, c, d = 4, 6, 3, 4
# func1_orig is the original loop-based version
%timeit func1_orig(a, b, c, d)
# 2.9 ms ± 110 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# func1 here is the vectorized version
%timeit func1(a, b, c, d)
# 210 µs ± 6.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

编辑:

请注意,以前的方法也不是全有或全无的事情。您可以选择仅对某些循环进行网格评估。例如,最里面的两个循环可以像这样被矢量化:

def func1(a, b, c, d):
    B = 0
    e = np.arange(a + b - 2).reshape((-1, 1))
    f = np.arange(c + d - 2)
    for ai in range(0, a):
        for bi in range(0, b):
            ei = e[:ai + bi]
            for ci in range(0, c):
                for di in range(0, d):
                    fi = f[:ci + di]
                    B += np.sum((2.) ** (ei - fi - ai - ci - di + 1) * (ei ** 2 - 2 * (ei * fi) - 7 * di) * scipy.special.factorial(ei))
    return B

这仍然有循环,但它确实避免了额外的计算,并且内存要求要低得多。哪个最好取决于我猜的输入的大小。在我的测试中,使用原始值 (4, 6, 3, 4) 这比原始函数还要慢;此外,对于这种情况,似乎在每个循环上为 eifi 创建新数组比在预先创建的切片上操作更快。但是,如果将输入乘以 4 (14, 24, 12, 16),那么这比原来的速度快得多(大约 x5),尽管仍然比完全矢量化的速度(大约 x3)慢。另一方面,我可以用这个(大约 5 分钟)计算输入缩放十(40、60、30、40)的值,但由于内存而不是前一个(我没有测试如何使用原始功能需要很长时间)。使用@numba.jit 有一点帮助,虽然不是很大(由于阶乘函数,不能使用nopython)。您可以根据输入的大小尝试对更多或更少的循环进行矢量化。

【讨论】:

  • 非常感谢您的出色回答。我没有考虑使用基于网格的方法,但它很好地解决了这个问题。我将对这两个选项进行试验,看看它们在提供的其他答案中的表现如何。我需要为 200,000 个 (a,b,c,d) 组合运行该函数,最多可达 50 个。我估计原始函数需要数月/数年才能完成。我需要解决的另一个问题是大整数的精度损失。也许 gmpy2 可能有用?
  • @Yeti 我刚刚添加了另一个不同的选项。预先计算昂贵的部分以避免重复计算并 JIT 编译其余部分似乎是这里的最佳选择。
  • @Yeti 关于精度,我对这个话题了解不够,无法给出建议(我的意思是,我知道添加很多数字会导致显着的精度损失,但我对方法了解不多解决这个问题)。我不知道将所有值存储在一个数组中并在最后将它们与np.sum(占用内存)相加是否会有所不同。在我做的实验中,这两种方法(累加和最终总和)似乎产生了非常接近的结果,即使不完全相同(不过我认为我一直使用 64 位浮点数)。
  • 我正在试验 JIT 编译选项;我一定很愚蠢,但我看不出exp_min 来自哪里,但exp_max 是有道理的。
  • @jdehesa 我在 Numba 中实现了因式分解,这带来了相当大的改进。您可以将此添加到您的答案中,然后我将删除我的(仅作为评论)
【解决方案2】:

使用这个cartesian_product 函数 您可以将嵌套循环转换为矩阵,然后您可以简单地以矢量化方式计算各自的嵌套 sigma:

In [37]: def nested_sig(args):
    ...:     base_prod = cartesian_product(*arrays)
    ...:     second_prod = cartesian_product(base_prod[:,:2].sum(1), base_prod[:,2:].sum(1))
    ...:     total = np.column_stack((base_prod, second_prod))
    ...:     # the items in each row denotes the following variables in order:
    ...:     # ai, bi, ci, di, ei, fi
    ...:     x = total[:, 4] - total[:, 5] - total[:, 0] - total[:, 2] - total[:, 3] + 1
    ...:     y = total[:, 4] - total[:, 5]
    ...:     result = np.power(2, x) * (np.power(total[:, 4], 2) - 2*y - 7*total[:, 3]) * np.math.factorial(total[:,4])
    ...:     return result

【讨论】:

    【解决方案3】:

    我在您的代码中看到了三个改进来源:

    • range(0,a)

    • 你在内部循环中做了很多工作

    • 您以随机方式对术语求和,存在较大条目丢失精度的风险。

    这里有一个版本(可能还不是很好)试图改进这一点。

    @numba.njit
    def func1o(a,b,c,d):
        "2**(ei-fi-ai-ci-di+1)*(ei**2-2*(ei*fi)-7*di)*ei!"                    
        POW=2.;                 SUM=0.;              
        L=[]
        for ai in arange(0.,a+1):
            for bi in range(0,b+1):
                for ci in range(0,c+1):
                    for di in range(0,d+1):
                        FACT=1.
                        for ei in arange(0,ai+bi+1):
                            for fi in range(0,ci+di+1):
                                L.append(POW*SUM*FACT)
                                POW /= 2
                                SUM -= 2*ei
                            POW *= 2    
                            SUM += 2*(ei-fi)+1
                            FACT *= ei+1
                        POW /=2
                        SUM -= 7*di
                    POW /= 2
            POW /= 2
        A=np.array(L)
        I=np.abs(A).argsort()
        return A[I].sum()    
    

    【讨论】:

      【解决方案4】:

      这只是对@jdehesa 答案的评论。

      如果 Numba 本身不支持某个功能,通常建议您自己实现它。在分解的情况下,这不是一项复杂的任务。

      代码

      import numpy as np
      import numba as nb
      
      @nb.njit()
      def factorial(a):
        res=1.
        for i in range(1,a+1):
          res*=i
        return res
      
      @nb.njit()
      def func1(a, b, c, d):
          B = 0.
      
          exp_min = 5 - (a + b + c + d)
          exp_max = b
          exp = 2. ** np.arange(exp_min, exp_max + 1)
      
          fact_e=np.empty(a + b - 2)
          for i in range(a + b - 2):
            fact_e[i]=factorial(i)
      
          for ai in range(0, a):
              for bi in range(0, b):
                  for ci in range(0, c):
                      for di in range(0, d):
                          for ei in range(0, ai + bi):
                              for fi in range(0, ci + di):
                                  B += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]
          return B
      

      并行版

      @nb.njit(parallel=True)
      def func_p(a_vec,b_vec,c_vec,d_vec):
        res=np.empty(a_vec.shape[0])
        for i in nb.prange(a_vec.shape[0]):
          res[i]=func1(a_vec[i], b_vec[i], c_vec[i], d_vec[i])
        return res
      

      示例

      a_vec=np.random.randint(low=2,high=10,size=1000000)
      b_vec=np.random.randint(low=2,high=10,size=1000000)
      c_vec=np.random.randint(low=2,high=10,size=1000000)
      d_vec=np.random.randint(low=2,high=10,size=1000000)
      
      res_2=func_p(a_vec,b_vec,c_vec,d_vec)
      

      单线程版本在您的示例中导致 5.6µs(第一次运行后)。

      并行版本几乎会导致另一个 Number_of_Cores 加速计算许多值。请记住,并行版本的编译开销更大(第一次调用超过 0.5 秒)。

      【讨论】:

      • 谢谢你,我正在调查它,但我收到了一个 Numba 错误,它(除其他外)说 This should not have happened, a problem has occurred in Numba's internals。您使用的是哪个版本的 Numba?我有 0.39.0。
      • @jdehesa 也是 0.39.0。 This should not have happened... 上面是什么?类似Untyped global name...cannot determine Numba type of...
      • 完整的错误消息是thisfunc1_2 是你的func1)。不知道那是什么,因为我后来更新了我的功能基本相同并且它有效(请参阅我的更新)。我还注意到fact_e 中的每个元素都可以从前一个元素中计算出来,而不是调用factorial(顺便说一句,不需要删除答案,因为它是一个有价值的贡献)。
      猜你喜欢
      • 2023-02-08
      • 1970-01-01
      • 2019-07-08
      • 2011-12-31
      • 2023-02-16
      • 2018-11-10
      • 1970-01-01
      • 1970-01-01
      • 2014-03-15
      相关资源
      最近更新 更多