Skip to main content
 首页 » 编程设计

python之numpy 函数的省时组合

2025年05月04日27mayingbao

我有两个 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?

请您参考如下方法:

Is there any way to do that accessing each element in the memory one time (thus make it time-efficient) without the need to write a custom 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