странное поведение numba при присваивании массиву - PullRequest
Новая
       6

странное поведение numba при присваивании массиву

0 голосов
/ 24 апреля 2018

У меня есть функция, которую я выполняю с помощью @jit (nopython = True).

Внутри него есть цикл, который делает кучу вещей, вычисляет корреляцию и затем назначает ее предварительно выделенному выходному массиву. И целевой массив, и корреляция имеют один и тот же тип (np.float32), но по какой-то причине назначение заставляет функцию занимать 100X дольше.

Чтобы сделать все еще более странным, если вместо моего значения корреляции вместо значения корреляции я назначу бессмысленное значение с плавающей точкой np.float32 (i * 1.01), функция будет работать с соответствующей скоростью.

Учитывая, что все имеют одинаковый тип, они оба должны работать с одинаковой скоростью, нет? 

corrs = np.zeros(a.shape[0], dtype=np.float32)

for i in range(lb, a.shape[0]):
    #a bunch of calculations happens here

    correl = np.float32(covar/(a_std*b_std))

    testval = np.float32(i*1.01)

    #doing this makes the function take FOREVER
    #corrs[i] = correl

    #but doing this runs very quickly, even though it is also a np.float32
    #corrs[i] = testval

вот пример для запуска. Я добавил аргумент «assign», который, если true, назначит то, что я хочу назначить, а если false, то назначит мое бесполезное значение теста. 

@jit(nopython=True)
def hist_corr_loop(a, b, lb = 1000, assign=True):


flb = np.float32(lb)

a_mu, b_mu = a[0], b[0]


for i in range(1, lb):
    a_mu+=a[i]
    b_mu+=b[i]



a_mu = a_mu/flb
b_mu = b_mu/flb


a_var, b_var = np.float32(0.0), np.float32(0.0)
for i in range(lb):
    a_var += np.square(a[i] - a_mu)
    b_var += np.square(b[i] - b_mu)

a_var = a_var/flb
b_var = b_var/flb


corrs = np.zeros(a.shape[0], dtype=np.float32)


for i in range(lb, a.shape[0]):

    #calculate new means and stdevs
    _a_mu = a_mu
    _b_mu = b_mu

    a_mu = _a_mu + (a[i] - a[i-lb])/flb
    b_mu = _b_mu + (b[i] - b[i-lb])/flb

    a_var += (a[i] - a[i-lb])*(a[i] - a_mu + a[i-lb] - _a_mu)/flb
    b_var += (b[i] - b[i-lb])*(b[i] - b_mu + b[i-lb] - _b_mu)/flb

    a_std = np.sqrt(a_var)#**0.5
    b_std = np.sqrt(b_var)#**0.5

    covar = np.float32(0.0)

    for j in range(i-lb+1,i+1):

        covar += (a[j] - a_mu)*(b[j] - b_mu)

    covar = covar/flb

    correl = np.float32(covar/(a_std*b_std))

    testval = np.float32(i*1.01)

    if assign:
        corrs[i] = correl

    else:
        corrs[i] = testval

return corrs

для запуска: 

n = 10000000
a = np.random.random(n)
b = np.random.random(n)

%timeit hist_corr_loop(a,b,1000, True)
%timeit hist_corr_loop(a,b, 1000, False)

Я получаю 

%timeit hist_corr_loop(a,b,1000, True)
10.5 s ± 52.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit hist_corr_loop(a,b, 1000, False)
220 ms ± 1.05 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

10 секунд против 220 мс.

...