首页 > 解决方案 > 使用 numpy 对带有条件代码的循环进行向量化

问题描述

我正在尝试将简单代码转换为矢量化器。我对这个概念很陌生。

我已经尝试使用该np.vectorize选项,但它会引发错误:

你能帮我解决这个问题吗?有没有其他更好的方法来矢量化它?


代码:

   def test(s,m):
      mean_matrix=s
      m2=s/m // forming an array
      result=[]
      for x in m2:
         return x/5



  s=np.random.random(50)
  m=np.random.random(100)
  temp=np.vectorize(test,otypes=[np.float],cache=False)
  out = temp(s,m)

标签: pythonnumpyloopsconditional-statementsvectorization

解决方案


首先将您的代码拆分为循环中使用的部分(可以矢量化)和其余部分(已经矢量化):

def kernel(s, x):
    i = 0
    initial_value = s[0].copy()
    while initial_value < x:
        i = i+1
        initial_value += s[i]
    return i

def test_split(s, m):
    mean_matrix = s.mean() # 5 usec
    product_m2 = mean_matrix * m # forming an array, 2 usec
    return [kernel(s, x) for x in product_m2] # 167 usec

test_split做同样的事情test,只是重构。我已经对每一行进行了注释,说明在那里花费了多少时间,以表明加快循环速度会有所作为。

如果你使用np.vectorize()它看起来像这样:

def test_vectorize(s, m):
    mean_matrix = s.mean() # 5 usec
    product_m2 = mean_matrix * m # forming an array, 2 usec
    kernel_vec = np.vectorize(lambda x: kernel(s, x)) # 2 usec
    return kernel_vec(product_m2) # 194 usec

这个比较慢!那是因为np.vectorize()不会让事情变得更快,它只会改变美学。它仍然在内部使用 Python 循环。最好忘记它的存在。

所以我们需要一种不同的方法,为此我们需要考虑做什么kernel。在我看来,它所做的是找到s总和小于的最长初始子序列xs.cumsum()将为我们提供所需的运行总数,但我们如何搜索小于 each xin的最大长度product_m2?我们不能使用循环,否则它会再次变慢。但是我们可以使用np.searchsorted(). 这里是:

def test_fast(s, m):
     mean_matrix = s.mean() # 5 usec
     product_m2 = mean_matrix * m # forming an array, 2 usec
     s_sum = s.cumsum() # 1 usec
     return s_sum.searchsorted(product_m2) # 2 usec

这样我们就有了 14 倍的加速。


推荐阅读