快速Numpy循环

时间:2016-06-13 15:14:11

标签: python numpy vectorization cython

如何优化此代码( 没有 矢量化,因为这会导致使用计算的语义,这通常远非重要) :

slow_lib.py:
import numpy as np

def foo():
    size = 200
    np.random.seed(1000031212)
    bar = np.random.rand(size, size)
    moo = np.zeros((size,size), dtype = np.float)
    for i in range(0,size):
        for j in range(0,size):
            val = bar[j]
            moo += np.outer(val, val)

关键是这种类型的循环经常与对某些向量运算有双重和的运算相对应。

这很慢:

>>t = timeit.timeit('foo()', 'from slow_lib import foo', number = 10)
>>print ("took: "+str(t))
took: 41.165681839

好的,那么让我们把它狠狠地加上它并添加类型注释,就像没有明天一样:

c_slow_lib.pyx:
import numpy as np
cimport numpy as np
import cython
@cython.boundscheck(False)
@cython.wraparound(False)

def foo():
    cdef int size = 200
    cdef int i,j
    np.random.seed(1000031212)
    cdef np.ndarray[np.double_t, ndim=2] bar = np.random.rand(size, size)
    cdef np.ndarray[np.double_t, ndim=2] moo = np.zeros((size,size), dtype = np.float)
    cdef np.ndarray[np.double_t, ndim=1] val
    for i in xrange(0,size):
        for j in xrange(0,size):
            val = bar[j]
            moo += np.outer(val, val)


>>t = timeit.timeit('foo()', 'from c_slow_lib import foo', number = 10)
>>print ("took: "+str(t))
took: 42.3104710579

...... ehr ......什么? Numba救援!

numba_slow_lib.py:
import numpy as np
from numba import jit

size = 200
np.random.seed(1000031212)

bar = np.random.rand(size, size)

@jit
def foo():
    bar = np.random.rand(size, size)
    moo = np.zeros((size,size), dtype = np.float)
    for i in range(0,size):
        for j in range(0,size):
            val = bar[j]
            moo += np.outer(val, val)

>>t = timeit.timeit('foo()', 'from numba_slow_lib import foo', number = 10)
>>print("took: "+str(t))
took: 40.7327859402

那么真的没有办法加快速度吗?重点是:

  • 如果我将内部循环转换为矢量化版本(构建一个代表内部循环的较大矩阵,然后在较大的矩阵上调用np.outer),我得到很多更快的代码。
  • 如果我在Matlab(R2016a)中实现类似的功能,由于JIT,这表现得相当不错。

4 个答案:

答案 0 :(得分:14)

以下是outer的代码:

def outer(a, b, out=None):    
    a = asarray(a)
    b = asarray(b)
    return multiply(a.ravel()[:, newaxis], b.ravel()[newaxis,:], out)

因此每次调用outer都会涉及许多python调用。那些最终调用编译代码来执行乘法。但是每个都会产生与数组大小无关的开销。

因此对outer的200(200 ** 2?)次调用将具有所有这些开销,而对所有200行的outer的一次调用具有一个开销集,然后是一次快速编译操作。

cythonnumba不会编译或以其他方式绕过outer中的Python代码。他们所能做的只是简化您编写的迭代代码 - 这并不会消耗太多时间。

没有深入细节,MATLAB jit必须能够用更快的代码替换'outer' - 它会重写迭代。但我对MATLAB的经验可以追溯到jit之前的时间。

对于使用cythonnumba的实际速度改进,您需要一直使用原始的numpy / python代码。或者更好地将你的精力集中在缓慢的内部碎片上。

用简化的版本替换outer可将运行时间缩短一半:

def foo1(N):
        size = N
        np.random.seed(1000031212)
        bar = np.random.rand(size, size)
        moo = np.zeros((size,size), dtype = np.float)
        for i in range(0,size):
                for j in range(0,size):
                        val = bar[j]
                        moo += val[:,None]*val   
        return moo

使用完整的N=200,你的函数每个循环需要17秒。如果我用pass替换内部两行(无计算),则每个循环的时间减少到3ms。换句话说,外循环机制不是一个很大的时间消费者,至少与许多outer()的调用没有比较。

答案 1 :(得分:9)

内存允许,你可以使用np.einsum以矢量化的方式执行那些繁重的计算,就像这样 -

moo = size*np.einsum('ij,ik->jk',bar,bar)

也可以使用np.tensordot -

moo = size*np.tensordot(bar,bar,axes=(0,0))

或只是np.dot -

moo = size*bar.T.dot(bar)

答案 2 :(得分:4)

Cython,Numba等的许多教程和演示使得这些工具看起来好像可以自动加速代码,但在实践中,通常情况并非如此:您需要稍微修改一下代码提取最佳表现。如果你已经实现了某种程度的矢量化,通常意味着写出所有的循环。 Numpy数组操作不是最优的原因包括:

  • 创建并循环了许多临时数组;
  • 如果数组很小,则每次调用开销很大;
  • 无法实现短路逻辑,因为数组是作为一个整体处理的;
  • 有时,无法使用数组表达式表示最优算法,并且您需要考虑时间复杂度较差的算法。

使用Numba或Cython不会优化这些问题!相反,这些工具允许您编写比普通Python快得多的循环代码。

此外,对于Numba,您应该了解"object mode" and "nopython mode"之间的区别。来自示例的紧密循环必须以nopython模式运行,以提供任何显着的加速。但是,numpy.outernot yet supported by Numba,导致函数在对象模式下编译。用jit(nopython=True)装饰,让这种情况抛出异常。

证明加速的例子确实可行:

import numpy as np
from numba import jit

@jit
def foo_nb(bar):
    size = bar.shape[0]
    moo = np.zeros((size, size))
    for i in range(0,size):
        for j in range(0,size):
            val = bar[j]
            moo += np.outer(val, val)
    return moo

@jit
def foo_nb2(bar):
    size = bar.shape[0]
    moo = np.zeros((size, size))
    for i in range(size):
        for j in range(size):
            for k in range(0,size):
                for l in range(0,size):
                    moo[k,l] += bar[j,k] * bar[j,l]
    return moo

size = 100
bar = np.random.rand(size, size)

np.allclose(foo_nb(bar), foo_nb2(bar))
# True

%timeit foo_nb(bar)
# 1 loop, best of 3: 816 ms per loop
%timeit foo_nb2(bar)
# 10 loops, best of 3: 176 ms per loop

答案 3 :(得分:-1)

您向我们展示的示例是一种效率低下的算法,因为您多次计算相同的外部产品。得到的时间复杂度为O(n ^ 4)。它可以减少到n ^ 3。

for i in range(0,size):
    val = bar[i]
    moo += size * np.outer(val, val)