在 Cython 中调用 numpy 函数确实会减慢速度

Question

我正在尝试在二维 numpy 数组中逐行调用 np.random.choice，无需替换。我正在使用 Cython 来提高速度。该代码的运行速度仅比纯 python 实现快 3 倍，这并不是一个很好的结果。瓶颈是 numpy 函数调用本身。当我把它注释掉，并且只为每一行提供一个静态结果，比如 [3, 2, 1, 0] 时，我得到了 1000 倍的加速（当然，它并没有做任何事情:)

我的问题：在调用 numpy 函数时我做错了什么导致它变得超级慢？从理论上讲，它是 C 与 C 对话，所以它应该很快。我查看了编译后的代码，对 numpy 函数的调用看起来很复杂，像 __Pyx_GOTREF 和 __Pyx_PyObject_GetAttrStr 这样的语句让我相信它在这个过程中使用了纯 Python（不好！！）。

我的代码：

# tag: numpy

import numpy as np

# compile-time info for numpy
cimport numpy as np
np.import_array()

# array dtypes
W_DTYPE = np.float
C_DTYPE = np.int

cdef int NUM_SELECTIONS = 4  # FIXME should be function kwarg

#compile-time dtypes
ctypedef np.float_t W_DTYPE_t
ctypedef np.int_t C_DTYPE_t


def allocate_choices(np.ndarray[W_DTYPE_t, ndim=2] round_weights,
                     np.ndarray[C_DTYPE_t, ndim=1] choice_labels):
    """
    For ea. row in  `round_weights` select NUM_SELECTIONS=4 items among
    corresponding `choice_labels`, without replacement, with corresponding
    probabilities in `round_weights`.

    Args:
        round_weights (np.ndarray): 2-d array of weights, w/
            size [n_rounds, n_choices]
        choice_labels (np.ndarray): 1-d array of choice labels,
            w/ size [n_choices]; choices must be *INTEGERS*

    Returns:
        choices (np.ndarray): selected items per round, w/ size
            [n_rounds, NUM_SELECTIONS]
    """

    assert round_weights.dtype == W_DTYPE
    assert choice_labels.dtype == C_DTYPE
    assert round_weights.shape[1] == choice_labels.shape[0]

    # initialize final choices array
    cdef int n_rows = round_weights.shape[0]
    cdef np.ndarray[C_DTYPE_t, ndim=2] choices = np.zeros([n_rows, NUM_SELECTIONS],
                                                          dtype=C_DTYPE)

    # Allocate choices, per round
    cdef int i, j
    cdef bint replace = False
    for i in range(n_rows):
        choices[i] = np.random.choice(choice_labels,
                                      NUM_SELECTIONS,
                                      replace,
                                      round_weights[i])

    return choices

Answer 1

在与一些人聊天并检查编译后的代码后对此进行更新：@DavidW 上面的评论说得很好：

“理论上它是 C 与 C 对话，所以它应该很快” - 不。不对。 cimport numpy 可以直接访问的 Numpy 的主要部分是 只是更快的数组索引。 Numpy 函数使用 正常的 Python 机制。它们最终可能在 C 中实现，但是 从 Cython 的角度来看，这并没有提供捷径。

所以这里的问题是，调用这个 Numpy 函数需要将输入转换回 python 对象，将它们传入，然后让 numpy 做它的事情。我不认为所有 Numpy 函数都是这种情况（从计时实验来看，其中一些我称之为工作很快），但很多都不是“Cythonized”。