python - 省时的numpy函数组合

我有两个3维的numpy数组A，B（大小〜（（1000，1000，3）->图像处理））以及基于元素的函数。

功能依次是：

import numpy as np
A  = A ** 3
A = np.maximum(A, 0.001)
C = np.divide(B, A)

由于操作这3个命令的功能是时间紧迫的瓶颈，我想问一下是否有一种方法可以通过对内存中的每个元素的单次访问来执行所有这些操作，即最快的性能。

我能找到的唯一组合是分隔部分，例如here或here或this one，这是特殊情况，因为爱因斯坦总和。

有没有一种方法可以一次访问内存中的每个元素（从而使它节省时间），而无需编写自定义ufunc？

最佳答案

有没有一种方法可以一次访问内存中的每个元素（从而使它节省时间），而无需编写自定义ufunc？


是的，这正是numexpr设计的目的。

import numpy as np
import numexpr as ne

def func1(A, B):
    A = A ** 3
    A = np.maximum(A, 0.001)
    return np.divide(B, A)

def func2(A, B):
    return ne.evaluate("B / where(A**3 > 0.001, A**3, 0.001)",
                       local_dict={'A':A,'B':B})

A, B = np.random.randn(2, 1000, 1000, 3)

print(np.allclose(func1(A, B), func2(A, B)))
# True

numexpr比原始代码提高了70倍：

In [1]: %%timeit A, B = np.random.randn(2, 1000, 1000, 3)
func1(A, B)
   ....:
1 loop, best of 3: 837 ms per loop

In [2]: %%timeit A, B = np.random.randn(2, 1000, 1000, 3)
func2(A, B)
   ....:
The slowest run took 8.87 times longer than the fastest. This could mean that an
intermediate result is being cached.
100 loops, best of 3: 11.5 ms per loop

部分原因是numexpr默认情况下将使用多个线程进行计算，但是即使使用单个线程，它仍然会破坏朴素的矢量化：

In [3]: ne.set_num_threads(1)
Out[3]: 8

In [4]: %%timeit A, B = np.random.randn(2, 1000, 1000, 3)
func2(A, B)
   ....:
10 loops, best of 3: 47.3 ms per loop

关于python - 省时的numpy函数组合，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/36958732/