Question

我们有一个使用掩码的矢量numpy get_pos_neg_bitwise 函数= [132 20 192] 和（500e3,4）的df.shape，我们想用numba加速。

from numba import jit
import numpy as np
from time import time

def get_pos_neg_bitwise(df, mask):
    """
    In [1]: print mask
    [132  20 192]

    In [1]: print df
    [[  1 162  97  41]
     [  0 136 135 171]
     ...,
     [  0 245  30  73]]

    """
    check = (np.bitwise_and(mask, df[:, 1:]) == mask).all(axis=1)
    pos = (df[:, 0] == 1) & check
    neg = (df[:, 0] == 0) & check
    pos = np.nonzero(pos)[0]
    neg = np.nonzero(neg)[0]
    return (pos, neg)

使用@morningsun的提示我们制作了这个numba版本：

@jit(nopython=True)
def numba_get_pos_neg_bitwise(df, mask):
    posneg = np.zeros((df.shape[0], 2))
    for idx in range(df.shape[0]):
        vandmask = np.bitwise_and(df[idx, 1:], mask)

        # numba fail with # if np.all(vandmask == mask):
        vandm_equal_m = 1
        for i, val in enumerate(vandmask):
            if val != mask[i]:
                vandm_equal_m = 0
                break
        if vandm_equal_m == 1:
            if df[idx, 0] == 1:
                posneg[idx, 0] = 1
            else:
                posneg[idx, 1] = 1
    pos = list(np.nonzero(posneg[:, 0])[0])
    neg = list(np.nonzero(posneg[:, 1])[0])
    return (pos, neg)

但它仍比numpy慢3倍（~0.06s Vs~0,02s）。

if __name__ == '__main__':

    df = np.array(np.random.randint(256, size=(int(500e3), 4)))
    df[:, 0] = np.random.randint(2, size=(1, df.shape[0]))  # set target to 0 or 1
    mask = np.array([132,  20, 192])

    start = time()
    pos, neg = get_pos_neg_bitwise(df, mask)
    msg = '==> pos, neg made; p={}, n={} in [{:.4} s] numpy'
    print msg.format(len(pos), len(neg), time() - start)

    start = time()
    msg = '==> pos, neg made; p={}, n={} in [{:.4} s] numba'
    pos, neg = numba_get_pos_neg_bitwise(df, mask)
    print msg.format(len(pos), len(neg), time() - start)
    start = time()
    pos, neg = numba_get_pos_neg_bitwise(df, mask)
    print msg.format(len(pos), len(neg), time() - start)

我错过了什么吗？

In [1]: %run numba_test2.py
==> pos, neg made; p=3852, n=3957 in [0.02306 s] numpy
==> pos, neg made; p=3852, n=3957 in [0.3492 s] numba
==> pos, neg made; p=3852, n=3957 in [0.06425 s] numba
In [1]:

Answer 1

尝试将调用移到循环外的np.bitwise_and，因为numba无法做任何事情来加快速度：

@jit(nopython=True)
def numba_get_pos_neg_bitwise(df, mask):
    posneg = np.zeros((df.shape[0], 2))
    vandmask = np.bitwise_and(df[:, 1:], mask)

    for idx in range(df.shape[0]):

        # numba fail with # if np.all(vandmask == mask):
        vandm_equal_m = 1
        for i, val in enumerate(vandmask[idx]):
            if val != mask[i]:
                vandm_equal_m = 0
                break
        if vandm_equal_m == 1:
            if df[idx, 0] == 1:
                posneg[idx, 0] = 1
            else:
                posneg[idx, 1] = 1
    pos = np.nonzero(posneg[:, 0])[0]
    neg = np.nonzero(posneg[:, 1])[0]
    return (pos, neg)

然后我得到了时间：

==> pos, neg made; p=3920, n=4023 in [0.02352 s] numpy
==> pos, neg made; p=3920, n=4023 in [0.2896 s] numba
==> pos, neg made; p=3920, n=4023 in [0.01539 s] numba

所以现在numba比numpy快一点。

此外，它没有产生太大的影响，但在原始函数中，您返回numpy数组，而在numba版本中，您将pos和neg转换为列表。

一般来说，我猜测函数调用由numpy函数控制，numba无法加速，而numpy版本的代码已经在使用快速矢量化例程。

<强>更新

您可以通过将enumerate调用和索引直接删除到数组而不是抓取切片来加快速度。同时将pos和neg拆分为单独的数组有助于避免在内存中沿着非连续轴切片：

@jit(nopython=True)
def numba_get_pos_neg_bitwise(df, mask):
    pos = np.zeros(df.shape[0])
    neg = np.zeros(df.shape[0])
    vandmask = np.bitwise_and(df[:, 1:], mask)

    for idx in range(df.shape[0]):

        # numba fail with # if np.all(vandmask == mask):
        vandm_equal_m = 1
        for i in xrange(vandmask.shape[1]):
            if vandmask[idx,i] != mask[i]:
                vandm_equal_m = 0
                break
        if vandm_equal_m == 1:
            if df[idx, 0] == 1:
                pos[idx] = 1
            else:
                neg[idx] = 1
    pos = np.nonzero(pos)[0]
    neg = np.nonzero(neg)[0]
    return pos, neg

ipython笔记本中的时间安排：

    %timeit pos1, neg1 = get_pos_neg_bitwise(df, mask)
    %timeit pos2, neg2 = numba_get_pos_neg_bitwise(df, mask)

    100 loops, best of 3: 18.2 ms per loop
    100 loops, best of 3: 7.89 ms per loop

Numba比numpy慢3倍

1 个答案: