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的官方文档中,才推荐大家使用nditerexternal_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 0column 1
row 099
row 132
row 259

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

  • 行优先方式
pos0pos1pos2pos3pos4pos5
935929
  • 列优先方式
pos0pos1pos2pos3pos4pos5
993259

事实上,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)

从上面的测试结果来看,nditerexternal_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(),定义如下:
    E M A t = α P r i c e t + ( 1 − α ) E M A t − 1 EMA_t = \alpha Price_t + (1-\alpha) EMA_{t-1} EMAt=αPricet+(1α)EMAt1看上去是一个典型的需要循环的函数,但通过寻找规律(等比数列),同样能够找到向量化的计算方法,从而实现效率的提升,具体可以参见这篇博文

好了,至少今天的实验得出两个结论:

  1. 尽可能地避免循环
  2. NumPy的nditer比python3.6的for-loop更慢,或者说,python3.6的for-loop似乎经过优化后比python2更快了,我在这篇博文里曾经做过python2.7和python3.6的循环效率对比
  • 6
    点赞
  • 18
    收藏
    觉得还不错? 一键收藏
  • 2
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值