Numba медленнее numpy.bitwise_and на булевых массивов

Question

Я пытаюсь Numba в этом фрагменте кода

from numba import jit 
import numpy as np 
from time import time 
db = np.array(np.random.randint(2, size=(400e3, 4)), dtype=bool) 
out = np.zeros((int(400e3), 1)) 

@jit() 
def check_mask(db, out, mask=[1, 0, 1]): 
    for idx, line in enumerate(db): 
     target, vector = line[0], line[1:] 
     if (mask == np.bitwise_and(mask, vector)).all(): 
      if target == 1: 
       out[idx] = 1 
    return out 

st = time() 
res = check_mask(db, out, [1, 0, 1]) 
print 'with jit: {:.4} sec'.format(time() - st) 

С Numba @jit() декоратора этот код работать медленнее!

• без JIT: 3,16 сек
• с JIT: 3,81 сек

просто, чтобы помочь лучше понять назначение этого кода:

db = np.array([   # out value for mask = [1, 0, 1] 
    # target, vector  # 
     [1,  1, 0, 1], # 1 
     [0,  1, 1, 1], # 0 (fit to mask but target == 0) 
     [0,  0, 1, 0], # 0 
     [1,  1, 0, 1], # 1 
     [0,  1, 1, 0], # 0 
     [1,  0, 0, 0], # 0 
     ]) 

Answer 1

Numba имеет два режима компиляции для jit: режим nopython и режим объекта. Режим Nopython (по умолчанию) поддерживает только ограниченный набор функций Python и Numpy, см. the docs for your version. Если функция jitted содержит неподдерживаемый код, Numba должен вернуться в режим объекта, что намного, намного медленнее.

Я не уверен, что режим objcet должен давать ускорение по сравнению с чистым Python, но в любом случае вы всегда будете использовать режим nopython. Для того, чтобы используется, что режим nopython, укажите nopython=True и придерживаться очень простой код (эмпирическое правило: выписывать все петли и использовать только скаляры и Numpy массивы):

@jit(nopython=True) 
def check_mask_2(db, out, mask=np.array([1, 0, 1])): 
    for idx in range(db.shape[0]): 
     if db[idx,0] != 1: 
      continue 
     check = 1 
     for j in range(db.shape[1]): 
      if mask[j] and not db[idx,j+1]: 
       check = 0 
       break 
     out[idx] = check 
    return out 

выписывая внутренний цикл явно имеет преимущество, которое мы можем вырваться из него, как только условие не удастся.

Тайминги:

%time _ = check_mask(db, out, np.array([1, 0, 1])) 
# Wall time: 1.91 s 
%time _ = check_mask_2(db, out, np.array([1, 0, 1])) 
# Wall time: 310 ms # slow because of compilation 
%time _ = check_mask_2(db, out, np.array([1, 0, 1])) 
# Wall time: 3 ms 

BTW, функция также легко векторизованные с Numpy, что дает приличную скорость:

def check_mask_vectorized(db, mask=[1, 0, 1]): 
    check = (db[:,1:] == mask).all(axis=1) 
    out = (db[:,0] == 1) & check 
    return out 

%time _ = check_mask_vectorized(db, [1, 0, 1]) 
# Wall time: 14 ms 

Answer 2

Я бы рекомендовал удаление вызова Numpy к array_equal из внутреннего цикла. numba не обязательно достаточно умна, чтобы превратить это в кусок встроенного C; и если он не заменит этот вызов, доминирующая стоимость вашей функции останется сопоставимой, что объяснит ваш результат.

В то время как numba может рассуждать о множестве конструкций numpy, это только код стиля C, действующий на массивах numpy, на которые можно положиться на ускорение.

Answer 3

В качестве альтернативы, вы можете попробовать pythran (отказ от ответственности: Я developper pythran).

С одной аннотации, он компилирует следующий код

#pythran export check_mask(bool[][], bool[]) 

import numpy as np 
def check_mask(db, out, mask=[1, 0, 1]): 
    for idx, line in enumerate(db): 
     target, vector = line[0], line[1:] 
     if (mask == np.bitwise_and(mask, vector)).all(): 
      if target == 1: 
       out[idx] = 1 
    return out 

с вызовом pythran check_call.py.

И согласно timeit, полученный родной модуль работает довольно быстро:

python -m timeit -s 'n=10e3 ; import numpy as np;db = np.array(np.random.randint(2, size=(n, 4)), dtype=bool); out = np.zeros(int(n),dtype=bool); from eq import check_mask' 'check_mask(db, out)' 

говорит мне версия CPython работает в 136ms в то время как pythran скомпилированные версия работает в 450us.