如何将此循环推入numpy？

Question

我有一个慢循环，我想通过将它推入 numpy.我花了几天时间玩这个代码却没有得到任何地方。甚至有可能，还是我缺少一个麻木的把戏？我可以做一些重构来提供帮助吗？

如您所见，我想将 mixin 相加为 xs 的偏移量。

import numpy as np

blocksize = 1000 # Chosen at runtime.
mixinsize = 100 # Chosen at runtime.
count = 10000 # Chosen at runtime.
xs = np.random.randint(0, blocksize + 1, count) # In practice this is data.
mixins = np.empty((count, mixinsize)) # In practice this is data.

# The slow part:
accumulator = np.zeros(blocksize + mixinsize)
for i in xrange(count):
    accumulator[xs[i]:xs[i] + mixinsize] += mixins[i]

对已接受答案的评论

• 要让 numba 工作，我首先必须确保我在任何地方都使用了适当的 numpy 类型，当然不是普通的 python 类型。这本身就是一个巨大的性能提升。
• Numba 改进了特定情况下的速度，因为计数变大了，这正是我想要的（也许我应该在问题中强调这一点）。
• 在我当前的测试用例中，不同的 numba 版本有明显不同的时间安排，但我认为 numba 是一个很有潜力的好主意，而且它正在积极开发中，所以我会坚持下去。
• 使用 anaconda 安装 numba，特别是它可以轻松地在 numba 版本之间切换。

Answer 1

从 numba.pydata.org 获取 Numba 0.11（还不是 0.12）。现在我们可以用 LLVM jit 编译这段代码了：

# plain NumPy version
import numpy as np

def foobar(mixinsize, count, xs, mixins, acc):
    for i in xrange(count):
        k = xs[i]
        acc[k:k + mixinsize] += mixins[i,:]


# LLVM compiled version
from numba import jit, void, int64, double
signature = void(int64,int64,int64[:],double[:,:],double[:])
foobar_jit = jit(signature)(foobar)

---

if __name__ == "__main__":   
    from time import clock

    blocksize = 1000 # Chosen at runtime.
    mixinsize = 100 # Chosen at runtime.
    count = 100000 # Chosen at runtime. 
    xs = np.random.randint(0, blocksize + 1, count)
    mixins = np.empty((count, mixinsize))
    acc = np.zeros(blocksize + mixinsize)

    t0 = clock()
    foobar(mixinsize, count, xs, mixins, acc)
    t1 = clock()
    print("elapsed time: %g ms" % (1000*(t1-t0),))

    t2 = clock()
    foobar_jit(mixinsize, count, xs, mixins, acc)
    t3 = clock()
    print("elapsed time with numba jit: %g ms" % (1000*(t3-t2),))

    print("speedup factor: %g" % ((t1-t0)/(t3-t2),))

---

$ python test_numba.py
elapsed time: 590.632 ms
elapsed time with numba jit: 12.31 ms
speedup factor: 47.9799

---

好的，只需增加三行 Python 代码，速度就几乎提高了 50 倍。

现在我们还可以使用 clang/LLVM 作为编译器来测试一个普通的 C 版本进行比较。

---

void foobar(long mixinsize, long count, 
    long *xs, double *mixins, double *accumulator)
{
    long i, j, k;
    double *cur, *acc;   
    for (i=0;i<count;i++) {
        acc = accumulator + xs[i];
        cur = mixins + i*mixinsize;
        for(j=0;j<mixinsize;j++) *acc++ += *cur++;
    }
}

---

from numpy.ctypeslib import ndpointer
import ctypes
so = ctypes.CDLL('plainc.so')
foobar_c = so.foobar
foobar_c.restype = None
foobar_c.argtypes = (
    ctypes.c_long,
    ctypes.c_long,
    ndpointer(dtype=np.int64, ndim=1),
    ndpointer(dtype=np.float64, ndim=2),
    ndpointer(dtype=np.float64, ndim=1)
)


t4 = clock()
foobar_c(mixinsize, count, xs, mixins, acc)
t5 = clock()
print("elapsed time with plain C: %g ms" % (1000*(t5-t4),))

---

$ CC -Ofast -shared -m64 -o plainc.so plainc.c
$ python test_numba.py
elapsed time: 599.136 ms
elapsed time with numba jit: 11.958 ms
speedup factor: 50.1034
elapsed time with plain C: 5.472 ms

---

因此，当使用 -Ofast 进行优化时，Numba 的速度大约是普通 C 版本的一半。相比之下，使用 -O2 的运行时间约为 8 毫秒。这意味着在这种情况下 Numba JIT 编译 Python 的速度大约是带有 -O2 优化标志的 C 的 75%。这对于仅仅增加三行 Python 代码来说还不错。

我们可以看一个普通的 Python 版本进行比较：

def foobar_py(mixinsize, count, xs, mixins, acc):
    for i in xrange(count):
        k = xs[i]
        for j in xrange(mixinsize):
            acc[j+k] += mixins[i][j]


# covert NumPy arrays to lists
_xs = map(int,xs)
_mixins = [map(float,mixins[i,:]) for i in xrange(count)]
_acc = map(float,acc)


t6 = clock()
foobar_py(mixinsize, count, _xs, _mixins, _acc)
t7 = clock()
print("elapsed time with plain Python: %g ms" % (1000*(t7-t6),))

此 Python 代码在 1775 毫秒内执行。因此，相对于普通 Python，我们可以使用 NumPy 获得大约 3 倍的加速，使用 Numba 可以获得 150 倍的加速，使用 C 和 -Ofast 可以获得 350 倍的加速。

警告来自 Donald Knuth，他将此归因于 CAR Hoare：“过早的优化是计算机编程中万恶之源。”要获得令人印象深刻的相对加速，沿着这条路线走下去的绝对加速只能让我们节省几毫秒的 CPU 时间。 我的时间真的值得从大量的劳动力中节省 CPU 吗？值得您的时间吗？自己决定。