Python学习笔记——Numpy - nditer循环的使用及效率研究

最新推荐文章于 2024-10-22 19:20:48 发布

置顶 QTEASY量化交易

最新推荐文章于 2024-10-22 19:20:48 发布

阅读量4.5k

点赞数 7

分类专栏： python 文章标签： python

本文链接：https://blog.csdn.net/Shepherdppz/article/details/104255289

版权

python 专栏收录该内容

53 篇文章 73 订阅

订阅专栏

Python学习笔记——Numpy-nditer循环的使用及效率研究

什么是`nditer`

nditer是NumPy库提供的“官方”ndarray的循环迭代器。

为什么要用`nditer`？

这的确是个好问题，大家都知道for-loop的速度是非常慢的，但凡在有可能的情况下，就应该使用numpy提供的向量化计算方法，能够带来百倍甚至千倍的速度提升。因此除非是万不得已，都不应该考虑使用循环的方法。举个简单的例子：

In [36]: def pow_loop(d): 
    ...:     res = np.zeros_like(d) 
    ...:     for i in range(len(d)): 
    ...:         res[i] = d[i] ** 2 
    ...:     return res 
    ...:

In [37]: def pow_vec(d): 
    ...:     return d ** 2 
    ...:   

In [38]: data = np.random.randint(10, size=(1000)) 

In [39]: np.allclose(pow_loop(data), pow_vec(data))
Out[39]: True

In [41]: %timeit pow_loop(data)
421 µs ± 2.25 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [42]: %timeit pow_vec(data)
1.09 µs ± 13.7 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [43]: 421 / 1.09
Out[43]: 386.23853211009174

上面两个函数都对同一个ndarray进行平方的操作，一个用了循环的方法，另一个是numpy的向量化方法，不出意料，向量化方法的速度是循环的386倍。

既然如此，为什么还需要讨论循环？那是因为在某些情况下很难将函数用向量化的方式实现。尽管绝大多数情况下总能找到向量化的办法。

总而言之，如果能够用向量化的方法实现的功能，一定要用向量化，只有在向量化没法使用的时候，才考虑用循环。我们现在就要来考察一下，如果我们迫不得已必须用循环的话，nditer能不能帮我们提升代码的效率。

`nditer`的基本用法

nditer是NumPy官方的循环迭代器，它的基本使用方法很简单，还是用前面的例子：

In [47]: def pow_nditer(d): 
    ...:     res = [] 
    ...:     it = np.nditer(d) 
    ...:     for item in it: 
    ...:         res.append(item ** 2) 
    ...:     return res 
    ...:              

In [48]: np.allclose(pow_loop(data), pow_nditer(data)) 
Out[48]: True

In [49]: %timeit pow_nditer(data) 
832 µs ± 9.32 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

尽管是“官方”迭代器，这速度可真够慢的。不过好在它胜在用法比较多，比如上面的例子，可以用这种方法来直接对矩阵中的元素进行修改，因此省去了构造list链表的功夫：

In [55]: def pow_nditer2(d): 
    ...:     res = np.empty_like(d) 
    ...:     it = np.nditer([d, res], op_flags=['readwrite'])
    ...:     for i, r in it: 
    ...:         r[...] = i ** 2 
    ...:     return res 
    ...:    

In [56]: np.allclose(pow_loop(data), pow_nditer2(data)) 
Out[56]: True

In [57]: %timeit pow_nditer2(data) 
1.13 ms ± 6.84 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

可惜，速度更慢了。因此在NumPy的官方文档中，才推荐大家使用nditer的external_loop模式，说这样的模式才能实现速度的提升，我们来看看external_loop模式是怎么回事：

In [58]: def pow_nditer_ext(d): 
    ...:     res = np.empty_like(d) 
    ...:     it = np.nditer([d, res], flags = ['external_loop'], op_flags=['readwrite'])
    ...:     for i, r in it: 
    ...:         r[...] = i ** 2 
    ...:     return res 
    ...: 
    
In [59]: np.allclose(pow_loop(data), pow_nditer_ext(data))
Out[59]: True

In [60]: %timeit pow_nditer_ext(data)
4.44 µs ± 34.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

啊哈！这样看上去好多了，这才像NumPy的样子嘛，可是这个external_loop是怎么回事呢？我们可以用打印的方式来测试一下，这次我们使用一个二维矩阵：

In [61]: d2d = np.random.randint(10, size = (5, 3))

In [62]: it = np.nditer(d2d)

In [63]: it2 = np.nditer(d2d, flags = ['external_loop']) 

In [65]: list(i for i in it)
Out[65]: 
[array(9),
 array(9),
 array(8),
 array(3),
 array(2),
 array(6),
 array(5),
 array(9),
 array(9),
 array(1),
 array(0),
 array(0),
 array(2),
 array(3),
 array(6)]

In [66]: list(i for i in it2)
Out[66]: [array([9, 9, 8, 3, 2, 6, 5, 9, 9, 1, 0, 0, 2, 3, 6])]

上面我们定义了一个五行三列的矩阵，分别用internal loop的方式和external_loop的方式打印出循环的结构，于是这样能看得很清楚，internal loop就跟正常的循环一样，在整个过程中循环15次，每次从矩阵中提取出一个元素，而external_loop则是一次性把所有的元素全部提取出来，放到“循环外”操作，这样只需要循环一次，所以才被称为“external loop”外部循环的原因吧。

可能有伙伴要问了，那么这个external_loop不就是不循环嘛？它有什么用？

在一种情况下external_loop很有用，那就是需要逐行或逐列循环的时候。下面我们看看具体怎么实现

用external loop实现逐行或逐列循环

假设我们需要对一个矩阵进行循环操作，但是这样的循环是逐行进行的，也就是说，需要把一个2D矩阵中的每一行（或列）逐行（或逐列）提取出来，进行操作，并按这个方式循环遍历整个矩阵的所有行（列）。

这种情况当然可以用data[i, :]的方式，但是如果用nditer的话怎么实现呢？原来，nditer的external_loop有两种提取元素的方式，这取决于循环的时候指定的数据提取方向。

在创建一个2D矩阵的时候，虽然我们的数据从逻辑结构上来说是M行N列的数据，但是在计算机内存中的存储方式仍然是线性存储的，占据MXN个单元的内存，那么既然存储结构和逻辑结构不一样，自然就存在两种不同的存放方法，一种是把元素按“行优先”的方式存储，另一种是按照“列优先”的方式存储，举个例子，下面的表格包含3行2列6个数字，逻辑上是一个二维矩阵：

	column 0	column 1
row 0	9	9
row 1	3	2
row 2	5	9

存储在内存中的两种方式如下

行优先方式

pos0	pos1	pos2	pos3	pos4	pos5
9	3	5	9	2	9

列优先方式

pos0	pos1	pos2	pos3	pos4	pos5
9	9	3	2	5	9

事实上，C语言使用的是第一种方式存储一个数组，而Fortran使用的是第二种方式。NumPy中在生成ndarray的时候可以指定使用何种方式存储数据，而当我们给出了external_loop关键字，提取数据的时候也可以指定提取数据的方式，numpy会从内存中提取最多的整块数据出来。

比如说上面的数组，如果用行优先的方式存储，我们如果用行优先的方法提取数据，那么就会得到整个矩阵作为一块提取到循环体中。但如果我们试图用列优先的方式提取数据呢？因为数据已经按行存储了，如果numpy试图强行从pos0和pos3中提取出第一行出来，这样效率会非常低，因而在这种情况下，numpy会首先提取pos0～pos2这一块连续的数据出来作为第一个循环，然后再提取第二块，于是就实现了逐列的循环。

可想而知，如果我们想要逐行循环，那么只需要把数组存储为列优先的方式，而按照行优先的方式来提取数据就可以了

下面我们就来测试一下，在numpy中指定数据存储方式需要用order参数，order='C'代表C方式（列优先） order = 'F'代表Fortran方式（行优先）：

In [70]: d2d 
Out[70]: 
array([[9, 9, 8],
       [3, 2, 6],
       [5, 9, 9],
       [1, 0, 0],
       [2, 3, 6]])

In [71]: d_c = np.array(d2d, order = 'C')

In [72]: d_f = np.array(d2d, order = 'F') 

In [73]: it_c = np.nditer(d_c, flags = ['external_loop'], order = 'F')

In [74]: it_f = np.nditer(d_f, flags = ['external_loop'], order = 'C') 

In [75]: list(i for i in it_c)
Out[75]: [array([9, 3, 5, 1, 2]), array([9, 2, 9, 0, 3]), array([8, 6, 9, 0, 6])]

In [76]: list(i for i in it_f) 
Out[76]: 
[array([9, 9, 8]),
 array([3, 2, 6]),
 array([5, 9, 9]),
 array([1, 0, 0]),
 array([2, 3, 6])]

从上面的例子可以看到，通过设置不同的存储方式，并且在执行external_loop循环时指定不同的数据提取方式，就可以实现逐行或逐列的循环。

既然我们很关注效率，那么自然下面需要对逐行逐列的循环效率进行一下对比：

In [78]: data = np.random.randint(10, size=(500, 300))
In [85]: data                                                                                                                     
Out[85]: 
array([[7, 9, 1, ..., 9, 4, 2],
       [3, 7, 5, ..., 2, 1, 3],
       [9, 1, 6, ..., 7, 3, 4],
       ...,
       [4, 2, 1, ..., 3, 2, 6],
       [8, 1, 4, ..., 4, 8, 6],
       [7, 2, 5, ..., 5, 1, 3]])

In [87]: def pow_loop(d): 
    ...:     res = np.zeros_like(d) 
    ...:     for i in range(500): 
    ...:         for j in range(500): 
    ...:             res[i,j]=d[i,j] ** 2 
    ...:     return res 
    ...: 

In [88]: def pow_loop_per_row(d): 
    ...:     res = np.zeros_like(d) 
    ...:     for i in range(500): 
    ...:         res[i,:] = d[i,:] ** 2 
    ...:     return res 
    ...: 

In [89]: def pow_nditer(d): 
    ...:     res = np.zeros_like(d) 
    ...:     it = np.nditer([d, res], op_flags=['readwrite']) 
    ...:     for i, r in it: 
    ...:         r[...] = i ** 2 
    ...:     return res 
    ...:

In [93]: def pow_nditer_ext(d):  
    ...:     d = np.array(d, order = 'C') 
    ...:     res = np.zeros_like(d, order = 'C') 
    ...:     it = np.nditer([d, res], flags = ['external_loop'], op_flags=['readwrite'], order='F') 
    ...:     for i, r in it: 
    ...:         r[...] = i ** 2 
    ...:     return res 
    ...:

In [94]: np.allclose(pow_loop(data), pow_nditer(data))
Out[94]: True

In [95]: np.allclose(pow_loop(data), pow_nditer(data))
Out[95]: Tru

In [96]: np.allclose(pow_loop(data), pow_loop_per_row(data))
Out[96]: True

In [97]: %timeit pow_loop(data)            
127 ms ± 3.25 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [98]: %timeit pow_nditer(data) 
278 ms ± 1.85 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [99]: %timeit pow_nditer_ext(data)                   
1.99 ms ± 78.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [100]: %timeit pow_loop_per_row(data)
1.1 ms ± 8.71 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

从上面的测试结果来看，nditer的external_loop的确能够大大缩短循环的时间。不过，话说回来，nditer迭代器的确“没什么用”，因为它的效率甚至还赶不上for-loop。

`nditer`的高效替代方案

终于到最后了。仍然还是那句老话，只要能用向量化计算的，千万别用循环，或许现在还可以加一句，就算要用循环，for-loop就可以了，别用nditer，效率比for-loop还低。

那么通常情况下，有哪些循环的替代方案呢？下面给出一些例子：

如果要将一个函数逐行应用到ndarray的每一行，在函数的参数相同的情况下，完全不需要使用external_loop或任何形式的loop，可以使用NumPy的函数np.apply_along_axis()，具体可以参见我的另一篇博文
那么如果需要对一个ndarray的每一行执行一个函数，但函数的参数不一样时，用np.apply_along_axis()就不行了，这时可以考虑使用map()函数，同样可以参考这篇博文
如果一个函数本身就是循环定义的，比如指数平滑移动平均EMA()，定义如下：
$EMA_t = \alpha Price_t + (1-\alpha) EMA_{t-1}$ 看上去是一个典型的需要循环的函数，但通过寻找规律（等比数列），同样能够找到向量化的计算方法，从而实现效率的提升，具体可以参见这篇博文