快速Numpy循环

Question

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

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

...... ......什么？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,这表现得相当好.

Answer 1

这是以下代码outer:

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

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

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

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

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

对于真正的速度提升与cython和numba你需要使用原始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().

Answer 2

在内存允许的情况下,您可以使用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)

Answer 3

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

• 创建并循环了许多临时数组；
• 如果数组很小，则每次调用的开销很大；
• 无法实现短路逻辑，因为数组是整体处理的；
• 有时无法使用数组表达式来表达最佳算法，而您只能使用时间复杂度较差的算法。

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

此外，特别是对于 Numba，您应该了解"object mode" 和 "nopython mode"之间的区别。您示例中的紧密循环必须在 nopython 模式下运行才能提供任何显着的加速。但是，numpy.outer是尚未被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