快速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

好吧，那么让我们将其cynotheize并添加类型注释，就像没有明天一样：

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_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的调用将有所有的开销，而对所有200行的outer的一个调用有一个开销集，然后是一个快速编译的操作。

cython和numba不要编译或绕过outer中的Python代码。他们所能做的就是简化你所写的迭代代码——这并不需要花费太多时间。

MatlabJIT必须能够用更快的代码替换“外部”代码，而不必详细说明，它可以重写迭代。但我对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函数，每个循环花费17s。如果我用pass（不计算）替换内部的两行，则每个循环的时间将减少到3ms。换句话说，外循环机制不是一个大的时间消耗器，至少与许多对outer()的调用相比不是这样。

Answer 2

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

• 大量的临时数组被创建和循环
• 如果数组很小，则每次调用的开销很大
• 短路逻辑不能实现，因为数组是整体处理的
• 有时，最优算法不能用数组表达式来表示，而您会选择时间复杂度更差的算法。

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

另外，对于Numba，您应该知道"object mode" and "nopython mode"之间的区别。示例中的紧循环必须在nopython模式下运行，以提供任何显著的加速。但是，numpy.outer是not 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

Answer 3

如果内存允许，可以使用^{}以矢量化的方式执行那些繁重的计算，例如-

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

也可以使用^{}-

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

或者干脆np.dot-

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