为什么 numba 没有提高速度？

Question

以下函数是优化问题的一部分。它非常简单，但经常被调用。如果速度能提高一点就好了。我尝试过 Numba，但似乎我必须用 FORTRAN 编写它：

\n\n

import numpy as np\nfrom numba import autojit\n# I am using numbapro: from numbapro import autojit\n\n# minimal dataset\nn_comp_glob = 2\nx_glob = np.random.rand(3*n_comp) \nqs_glob = np.array([100.])\ncp_glob =np.tile(1.,n_comp)\ncs_glob = np.array( [100.])\n\ndef get_denom(n_comp,qs,x,cp,cs_f):\n    k = x[0:n_comp]\n    sigma = x[n_comp:2*n_comp]\n    z = x[2*n_comp:3*n_comp]\n    # calculates the denominator in Equ 14a - 14c (Brooks & Cramer 1992)\n    a = 0.0 \n\n    for i in range(n_comp):\n         a += (sigma[i] + z[i])*( k[i]*(qs/cs)**(z[i]-1) )*cp[i]\n\n    return denom\n\nget_denom_jit=autojit(get_denom)\n\nimport timeit\n%timeit get_denom(n_comp_glob,qs_glob,x_glob,cp,cs_glob)\n10000 loops, best of 3: 22.9 \xc2\xb5s per loop\n%timeit get_denom_jit(n_comp_glob,qs_glob,x_glob,cp_glob,cs_glob)\n10000 loops, best of 3: 27.9 \xc2\xb5s per loop\n

\n\n

如果我增加组件的数量（n_comp），numba 仍然不比 python 快。为什么 ？

\n\n

编辑：

\n\n

我尝试了从numba 文档中截取的代码：

\n\n

from numba import *\nimport numpy as np\n\nmu = 0.1\nLx, Ly = 101, 101\nN = 1000\n\n\ndef diffuse_loops(iter_num):\n    u = np.zeros((Lx, Ly), dtype=np.float64)\n    temp_u = np.zeros_like(u)\n    temp_u[Lx / 2, Ly / 2] = 1000.0\n\n    for n in range(iter_num):\n        for i in range(1, Lx - 1):\n            for j in range(1, Ly - 1):\n                u[i, j] = mu * (temp_u[i + 1, j] + temp_u[i - 1, j] +\n                            temp_u[i, j + 1] + temp_u[i, j - 1] -\n                            4 * temp_u[i, j])\n\n    temp = u\n    u = temp_u\n    temp_u = temp\n\nreturn u\n\n\ndef diffuse_array_expressions(iter_num):\n    u = np.zeros((Lx, Ly), dtype=np.float64)\n    temp_u = np.zeros_like(u)\n    temp_u[Lx / 2, Ly / 2] = 1000.0\n\n    for i in range(iter_num):\n    u[1:-1, 1:-1] = mu * (temp_u[2:, 1:-1] + temp_u[:-2, 1:-1] +\n                          temp_u[1:-1, 2:] + temp_u[1:-1, :-2] -\n                          4 * temp_u[1:-1, 1:-1])\n\n    temp = u\n    u = temp_u\n    temp_u = temp\n\nreturn u\n\ndiffuse_array_expressions_jit = autojit(diffuse_array_expressions)\ndiffuse_loops_jit = autojit(diffuse_loops)\n

\n\n

调用函数：

\n\n

%timeit diffuse_array_expressions(100)\n100 loops, best of 3: 13.9 ms per loop\n\n%timeit diffuse_array_expressions_jit(100)\n100 loops, best of 3: 14.8 ms per loop\n\n%timeit diffuse_loops(100)\n1 loops, best of 3: 1.88 s per loop\n\n%timeit diffuse_loops_jit(100)\n1000 loops, best of 3: 1.87 ms per loop\n

\n\n

所以看起来，当代码在 Python 中正确实现时，numba 就失去了它的力量。上面的函数当然可以在没有循环的情况下实现：

\n\n

def get_denom_vec(n_comp,qs,x,cp,cs_f):\n    k = x[0:n_comp]\n    sigma = x[n_comp:2*n_comp]\n    z = x[2*n_comp:3*n_comp]\n    # calculates the denominator in Equ 14a - 14c (Brooks & Cramer 1992)\n    a = 0.0 \n\n    a = (sigma + z)*( k*(qs/cs)**(z-1) )*cp\n    denom = np.sum(a) + cs\n    return denom\n\nget_denom_vec_jit = autojit(get_denom_vec)\n\n%timeit get_denom_vec_jit(n_comp_glob,qs_glob,x_glob,cp_glob,cs_glob)\n1000 loops, best of 3: 223 \xc2\xb5s per loop\n\n%timeit get_denom_vec_jit(n_comp_glob,qs_glob,x_glob,cp_glob,cs_glob)\n1000 loops, best of 3: 245 \xc2\xb5s per loop\n

\n

Answer 1

numba 的最酷之处在于您可以自己找到它！

使用

@autojit(nopython=True)

或者

numba --annotate myfile.py

并迭代修复无法jitted的代码。