NumPy是Python中众多科学软件包的基础。它提供了一个特殊的数据类型ndarray,其在向量计算上做了优化。这个对象是科学数值计算中大多数算法的核心。

相比于原生的Python,利用NumPy数组可以获得显著的性能加速,尤其是当你的计算遵循单指令多数据流(SIMD)范式时。然而,利用NumPy也有可能有意无意地写出未优化的代码。

在这篇文章中,我们将看到一些技巧,这些技巧可以帮助你编写高效的NumPy代码。我们首先看一下如何避免不必要的数组拷贝,以节省时间和内存。因此,我们将需要深入NumPy的内部。

学习避免不必要的数据拷贝

NumPy数组计算可能涉及到内存块之间的内部拷贝。有时会有不必要的拷贝,此时应该避免。相应地这里有一些技巧,可以帮助你优化你的代码。

import numpy as np

查看数组的内存地址

1. 查看静默数组拷贝的第一步是在内存中找到数组的地址。下边的函数就是做这个的:

def id(x):
# This function returns the memory
# block address of an array.
return x.__array_interface__['data'][0]

2. 有时你可能需要复制一个数组,例如你需要在操作一个数组时,内存中仍然保留其原始副本。

a = np.zeros(10); aid = id(a); aid
71211328
b = a.copy(); id(b) == aid
False

具有相同数据地址(比如id函数的返回值)的两个数组,共享底层数据缓冲区。然而,共享底层数据缓冲区的数组,只有当它们具有相同的偏移量(意味着它们的第一个元素相同)时,才具有相同的数据地址。共享数据缓冲区,但偏移量不同的两个数组,在内存地址上有细微的差别,正如下边的例子所展示的那样:

id(a), id(a[1:])
(71211328, 71211336)

在这篇文章中,我们将确保函数用到的数组具有相同的偏移量。下边是一个判断两个数组是否共享相同数据的更可靠的方案:

def get_data_base(arr):
"""For a given Numpy array, finds the base array that "owns" the actual data."""
base = arr
while isinstance(base.base, np.ndarray):
base = base.base
return base
def arrays_share_data(x, y):
return get_data_base(x) is get_data_base(y)
print(arrays_share_data(a,a.copy()), arrays_share_data(a,a[1:]))
False True

感谢Michael Droettboom指出这种更精确的方法,提出这个替代方案。

就地操作和隐式拷贝操作

3. 数组计算包括就地操作(下面第一个例子:数组修改)或隐式拷贝操作(第二个例子:创建一个新的数组)。