numpy 函数的高效组合答案

【问题标题】：Time-efficient combination of numpy functionsnumpy 函数的高效组合
【发布时间】：2016-08-25 19:39:01
【问题描述】：

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

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

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

有没有什么方法可以在不需要编写自定义 ufunc 的情况下一次性访问内存中的每个元素（从而使其节省时间）？

【问题讨论】：

数据类型允许，使用第三个链接，做np.einsum('ijk,ijk,ijk->ijk',A,A,A)模拟A***3？这应该是非常有效的。

标签： python performance numpy

【解决方案1】：

有没有什么方法可以在不需要编写自定义 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

【讨论】：

这个加速是巨大的！我最后确认了类似的数字。
确实令人印象深刻。谢谢！

【解决方案2】：

坦率地说，考虑到相关的加速数字，@ali_m's solution 中列出的 numexpr 单线看起来像是要走的路。在这篇文章中列出的将内容保存在 NumPy 中是另一种建议。

让我们一次一个地计算这些指令的时间，看看是否有任何瓶颈 -

In [108]: # Random input arrays
     ...: A = np.random.rand(1000,1000,3)
     ...: B = np.random.rand(1000,1000,3)
     ...: 

In [109]: %timeit A**3
1 loops, best of 3: 442 ms per loop

In [110]: A  = A ** 3

In [111]: %timeit np.maximum(A, 0.001)
100 loops, best of 3: 16.4 ms per loop

In [112]: A = np.maximum(A, 0.001)

In [113]: %timeit np.divide(B, A)
10 loops, best of 3: 19.7 ms per loop

因此，功率计算似乎占据了总运行时间的很大一部分。

让我们在那里介绍np.einsum，但请注意所涉及的数据类型。

In [114]: # Random input arrays
     ...: A = np.random.rand(1000,1000,3)
     ...: B = np.random.rand(1000,1000,3)
     ...: 

In [115]: %timeit A**3
1 loops, best of 3: 442 ms per loop

In [116]: %timeit np.einsum('ijk,ijk,ijk->ijk',A,A,A)
10 loops, best of 3: 28.3 ms per loop

In [117]: np.allclose(A**3,np.einsum('ijk,ijk,ijk->ijk',A,A,A))
Out[117]: True

这是一个很好的加速。

【讨论】：