Проблемы с ускорением функций с помощью numba JIT

Я новичок в jit numba. Для личного проекта мне нужно ускорить функции, похожие на те, что будут показаны ниже, но разные для целей написания автономных примеров.

import numpy as np
from numba import jit, autojit, double, float64, float32, void

def f(n):
    k=0.
    for i in range(n):
        for j in range(n):
            k+= i+j

def f_with_return(n):
    k=0.
    for i in range(n):
        for j in range(n):
            k+= i+j
    return k

def f_with_arange(n):
    k=0.
    for i in np.arange(n):
        for j in np.arange(n):
            k+= i+j

def f_with_arange_and_return(n):
    k=0.
    for i in np.arange(n):
        for j in np.arange(n):
            k+= i+j  


#jit decorators
jit_f = jit(void(int32))(f)
jit_f_with_return = jit(int32(int32))(f_with_return)
jit_f_with_arange = jit(void(double))(f_with_arange)
jit_f_with_arange_and_return = jit(double(double))(f_with_arange_and_return)

И ориентиры:

%timeit f(1000)
%timeit jit_f(1000)

10 циклов, лучшее из 3: 73,9 мс на цикл / 1000000 циклов, лучшее из 3: 212 нс на цикл

%timeit f_with_return(1000)
%timeit jit_f_with_return(1000)

10 циклов, лучшее из 3: 74,9 мс на цикл / 1000000 циклов, лучшее из 3: 220 нс на цикл

Я не понимаю этих двух:

%timeit f_with_arange(1000.0)
%timeit jit_f_with_arange(1000.0)

10 циклов, лучший из 3: 175 мс на цикл / 1 цикл, лучший из 3: 167 мс на цикл

%timeit f_with_arange_with_return(1000.0)
%timeit jit_f_with_arange_with_return(1000.0)

10 циклов, лучший из 3: 174 мс на цикл / 1 цикл, лучший из 3: 172 мс на цикл

Я думаю, что я не даю функции jit правильные типы для вывода и ввода? Просто потому, что цикл for теперь работает с numpy.arange, а не с простым диапазоном, я не могу заставить jit сделать его быстрее. В чем здесь проблема?


python numpy jit numba
person Mathusalem    schedule 06.03.2015    source источник

Ответы (1)


arrow_upward
7
arrow_downward

Просто numba не знает, как преобразовать np.arange в собственный цикл низкого уровня, поэтому по умолчанию он использует слой объектов, который намного медленнее и обычно имеет ту же скорость, что и чистый python.

Хороший трюк — передать аргумент ключевого слова nopython=True в jit, чтобы посмотреть, сможет ли он скомпилировать все, не прибегая к объектному режиму:

import numpy as np
import numba as nb

def f_with_return(n):
    k=0.
    for i in range(n):
        for j in range(n):
            k+= i+j
    return k

jit_f_with_return = nb.jit()(f_with_return)
jit_f_with_return_nopython = nb.jit(nopython=True)(f_with_return)

%timeit f_with_return(1000)
%timeit jit_f_with_return(1000)
%timeit jit_f_with_return_nopython(1000)

Последние два имеют одинаковую скорость на моей машине и намного быстрее, чем необработанный код. Два примера, о которых у вас были вопросы, вызовут ошибку с nopython=True, так как на данный момент он не может скомпилировать np.arange.

Дополнительные сведения см. в следующих разделах:

http://numba.pydata.org/numba-doc/0.17.0/user/troubleshoot.html#the-compiled-code-is-too-slow

и для списка поддерживаемых функций numpy с указанием того, что поддерживается и не поддерживается в режиме nopython:

http://numba.pydata.org/numba-doc/0.17.0/reference/numpysupported.html

person JoshAdel    schedule 06.03.2015