NumPy是Python中众多科学软件包的基础。它提供了一个特殊的数据类型ndarray,其在向量计算上做了优化。这个对象是科学数值计算中大多数算法的核心。
相比于原生的Python,利用NumPy数组可以获得显著的性能加速,尤其是当你的计算遵循单指令多数据流(SIMD)范式时。然而,利用NumPy也有可能有意无意地写出未优化的代码。
在这篇文章中,我们将看到一些技巧,这些技巧可以帮助你编写高效的NumPy代码。我们首先看一下如何避免不必要的数组拷贝,以节省时间和内存。因此,我们将需要深入NumPy的内部。
学习避免不必要的数据拷贝
NumPy数组计算可能涉及到内存块之间的内部拷贝。有时会有不必要的拷贝,此时应该避免。相应地这里有一些技巧,可以帮助你优化你的代码。
1 |
import numpy as np |
查看数组的内存地址
1. 查看静默数组拷贝的第一步是在内存中找到数组的地址。下边的函数就是做这个的:
1 2 3 4 |
def id(x): # This function returns the memory # block address of an array. return x.__array_interface__['data'][0] |
2. 有时你可能需要复制一个数组,例如你需要在操作一个数组时,内存中仍然保留其原始副本。
1 2 3 4 |
a = np.zeros(10); aid = id(a); aid 71211328 b = a.copy(); id(b) == aid False |
具有相同数据地址(比如id函数的返回值)的两个数组,共享底层数据缓冲区。然而,共享底层数据缓冲区的数组,只有当它们具有相同的偏移量(意味着它们的第一个元素相同)时,才具有相同的数据地址。共享数据缓冲区,但偏移量不同的两个数组,在内存地址上有细微的差别,正如下边的例子所展示的那样:
1 2 |
id(a), id(a[1:]) (71211328, 71211336) |
在这篇文章中,我们将确保函数用到的数组具有相同的偏移量。下边是一个判断两个数组是否共享相同数据的更可靠的方案:
1 2 3 4 5 6 7 8 9 10 11 12 |
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. 数组计算包括就地操作(下面第一个例子:数组修改)或隐式拷贝操作(第二个例子:创建一个新的数组)。
1 2 3 4 5 |
a *= 2; id(a) == aid True c = a * 2; id(c) == aid False |
一定要选择真正需要的操作类型。隐式拷贝操作很明显很慢,如下所示:
1 2 3 4 5 6 7 8 9 |
%%timeit a = np.zeros(10000000) a *= 2 10 loops, best of 3: 19.2 ms per loop %%timeit a = np.zeros(10000000) b = a * 2 10 loops, best of 3: 42.6 ms per loop |
4. 重塑一个数组可能涉及到拷贝操作,也可能涉及不到。原因将在下面解释。例如,重塑一个二维矩阵不涉及拷贝操作,除非它被转置(或更一般的非连续操作):
1 2 |
a = np.zeros((10, 10)); aid = id(a); aid 53423728 |
重塑一个数组,同时保留其顺序,并不触发拷贝操作。
1 2 |
b = a.reshape((1, -1)); id(b) == aid True |
转置一个数组会改变其顺序,所以这种重塑会触发拷贝操作。
1 2 |
c = a.T.reshape((1, -1)); id(c) == aid False |
因此,后边的指令比前边的指令明显要慢。
5. 数组的flatten和revel方法将数组变为一个一维向量(铺平数组)。flatten方法总是返回一个拷贝后的副本,而revel方法只有当有必要时才返回一个拷贝后的副本(所以该方法要快得多,尤其是在大数组上进行操作时)。
1 2 3 4 5 6 7 8 9 10 11 12 13 |
d = a.flatten(); id(d) == aid False e = a.ravel(); id(e) == aid True %timeit a.flatten() 1000000 loops, best of 3: 881 ns per loop %timeit a.ravel() 1000000 loops, best of 3: 294 ns per loop |
广播规则
6. 广播规则允许你在形状不同但却兼容的数组上进行计算。换句话说,你并不总是需要重塑或铺平数组,使它们的形状匹配。下面的例子说明了两个向量之间进行矢量积的两个方法:第一个方法涉及到数组的变形操作,第二个方法涉及到广播规则。显然第二个方法是要快得多。
1 2 3 4 5 6 7 8 9 10 11 12 13 |
n = 1000 a = np.arange(n) ac = a[:, np.newaxis] ar = a[np.newaxis, :] %timeit np.tile(ac, (1, n)) * np.tile(ar, (n, 1)) 100 loops, best of 3: 10 ms per loop %timeit ar * ac 100 loops, best of 3: 2.36 ms per loop |
在NumPy数组上进行高效的选择
NumPy提供了多种数组分片的方式。数组视图涉及到一个数组的原始数据缓冲区,但具有不同的偏移量,形状和步长。NumPy只允许等步长选择(即线性分隔索引)。NumPy还提供沿一个轴进行任意选择的特定功能。最后,花式索引(fancy indexing)是最一般的选择方法,但正如我们将要在文章中看到的那样,它同时也是最慢的。如果可能,我们应该选择更快的替代方法。
1. 创建一个具有很多行的数组。我们将沿第一维选择该数组的分片。
1 2 |
n, d = 100000, 100 a = np.random.random_sample((n, |