过滤 NumPy 数组

Question

假设我有一个 NumPy 数组 arr，我想对其进行逐元素过滤，例如我只想获得低于某个阈值 k.

的值

有几种方法，例如：

使用生成器：np.fromiter((x for x in arr if x < k), dtype=arr.dtype)
使用布尔掩码切片：arr[arr < k]
使用np.where()：arr[np.where(arr < k)]
使用np.nonzero()：arr[np.nonzero(arr < k)]
使用基于 Cython 的自定义实现
使用基于 Numba 的自定义实现

哪个最快？内存效率如何？

（编辑：添加 np.nonzero() 基于@ShadowRanger 评论）

Answer 1

总结

以下测试旨在提供对不同方法的一些见解，应持保留态度。

这里测试的不是一般的过滤，而是只是应用一个阈值，它的显着特点是计算条件非常快。如果条件意味着昂贵的计算，将获得非常不同的结果。

基本上 2-pass 加速（使用 Numba 或 Cython - 只要您事先知道类型）将是最快的并且内存效率更高，除了非常大的输入，单次通过Numba/Cython 更快（以使用更大的临时内存为代价）。使用 np.where()/np.nonzero()，而不是直接使用掩码，可能会导致计算速度稍微加快，并且通常不会造成伤害（除了可能占用较大的临时内存），只要输出的大小是小于输入的 50%。 np.fromiter() 方法 慢得多 ，但不会产生大型临时对象。

定义

使用生成器：

def filter_fromiter(arr, k):
    return np.fromiter((x for x in arr if x < k), dtype=arr.dtype)

使用布尔掩码切片：

def filter_mask(arr, k):
    return arr[arr < k]

使用np.where():

def filter_where(arr, k):
    return arr[np.where(arr < k)]

使用np.nonzero()

def filter_nonzero(arr, k):
    return arr[np.nonzero(arr < k)]

使用基于 Cython 的自定义实现：
- 单程filter_cy()
- 两遍filter2_cy()

%%cython -c-O3 -c-march=native -a
#cython: language_level=3, boundscheck=False, wraparound=False, initializedcheck=False, cdivision=True, infer_types=True


cimport numpy as cnp
cimport cython as ccy

import numpy as np
import cython as cy


cdef long NUM = 1048576
cdef long MAX_VAL = 1048576
cdef long K = 1048576 // 2


cdef int smaller_than_cy(long x, long k=K):
    return x < k


cdef size_t _filter_cy(long[:] arr, long[:] result, size_t size, long k):
    cdef size_t j = 0
    for i in range(size):
        if smaller_than_cy(arr[i]):
            result[j] = arr[i]
            j += 1
    return j


cpdef filter_cy(arr, k):
    result = np.empty_like(arr)
    new_size = _filter_cy(arr, result, arr.size, k)
    result.resize(new_size)
    return result


cdef size_t _filtered_size(long[:] arr, size_t size, long k):
    cdef size_t j = 0
    for i in range(size):
        if smaller_than_cy(arr[i]):
            j += 1
    return j


cpdef filter2_cy(arr, k):
    cdef size_t new_size = _filtered_size(arr, arr.size, k)
    result = np.empty(new_size, dtype=arr.dtype)
    new_size = _filter_cy(arr, result, arr.size, k)
    return result

import functools


filter_np_cy = functools.partial(filter_cy, k=K)
filter_np_cy.__name__ = 'filter_np_cy'


filter2_np_cy = functools.partial(filter2_cy, k=K)
filter2_np_cy.__name__ = 'filter2_np_cy'

使用基于 Numba 的自定义实现
- 单程filter_np_nb()
- 两遍filter2_np_nb()

import numba as nb
import functools


@nb.njit
def filter_func(x, k):
    return x < k


@nb.njit
def filter_nb(arr, result, k):
    j = 0
    for i in range(arr.size):
        if filter_func(arr[i], k):
            result[j] = arr[i]
            j += 1
    return j


def filter_np_nb(arr, k=K):
    result = np.empty_like(arr)
    j = filter_nb(arr, result, k)
    result.resize(j, refcheck=False)
    return result


@nb.njit
def filter2_nb(arr, k):
    j = 0
    for i in range(arr.size):
        if filter_func(arr[i], k):
            j += 1
    result = np.empty(j, dtype=arr.dtype)
    j = 0
    for i in range(arr.size):
        if filter_func(arr[i], k):
            result[j] = arr[i]
            j += 1
    return result


filter2_np_nb = functools.partial(filter2_nb, k=K)
filter2_np_nb.__name__ = 'filter2_np_nb'

计时基准

基于生成器的 filter_fromiter() 方法比其他方法慢得多（大约慢 2 个数量级，因此在图表中被省略）。

时间将取决于输入数组大小和过滤项目的百分比。

作为输入大小的函数

第一张图将时序作为输入大小的函数（约 50% 过滤掉的元素）：

一般来说，基于 Numba 的方法始终是最快的，紧随其后的是 Cython 方法。在它们中，两次通过方法通常是最快的，除了非常大的输入，在这种情况下，单次通过方法往往会接管。在 NumPy 中，基于 np.where() 和基于 np.nonzero() 的方法基本相同（除了非常小的输入，np.nonzero() 似乎稍微慢一些），并且它们都比布尔掩码切片，除了非常小的输入（低于 ~100 个元素），布尔掩码切片速度更快。此外，对于非常小的输入，基于 Cython 的解决方案比基于 NumPy 的解决方案慢。

作为填充函数

第二张图将时间作为项目通过过滤器的函数（对于约 100 万个元素的固定输入大小）：

第一个观察结果是，所有方法在接近 ~50% 填充时速度最慢，填充更少或更多时速度更快，并且在无填充时速度最快（过滤掉的值百分比最高，通过百分比最低通过图中 x 轴所示的值）。同样，Numba 和 Cython 版本通常都比基于 NumPy 的对应版本更快，Numba 几乎总是最快的，而 Cython 在图表最右边的部分胜过 Numba。两次通过的方法对于较大的填充值具有增加的边际速度增益，直到大约。 50%，之后单程接管速度领奖台。在 NumPy 中，基于 np.where() 和基于 np.nonzero() 的方法再次基本相同。比较基于 NumPy 的解决方案时，np.where()/np.nonzero() 解决方案在填充低于 ~60% 时优于布尔掩码切片，此后布尔掩码切片变得最快。

（完整代码可用 here）

内存注意事项

基于生成器的 filter_fromiter() 方法只需要最少的临时存储，与输入的大小无关。内存方面，这是最有效的方法。具有相似内存效率的是 Cython / Numba 两遍方法，因为输出的大小是在第一遍期间确定的。

在内存方面，Cython 和 Numba 的单通道解决方案都需要输入大小的临时数组。因此，与两次遍历或基于生成器的遍历相比，它们的内存效率不是很高。然而，与掩码相比，它们具有相似的渐近临时内存占用空间，但常数项通常大于掩码。

布尔掩码切片解决方案需要输入大小但类型为 bool 的临时数组，在 NumPy 中为 1 位，因此这比 NumPy 的默认大小小约 64 倍典型 64 位系统上的数组。

基于np.nonzero()/np.where()的解决方案与第一步（在np.nonzero()/np.where()内）的布尔掩码切片具有相同的要求，它被转换到第二步（np.nonzero()/np.where() 的输出）中的一系列 int（在 64 位系统上通常是 int64）。因此，第二步具有可变的内存要求，具体取决于过滤元素的数量。

备注

生成器方法在指定不同的过滤条件时也是最灵活的
Cython 解决方案需要指定数据类型以提高速度，或者为多类型调度付出额外的努力
对于Numba和Cython，过滤条件都可以指定为通用函数（因此不需要硬编码），但必须在各自的环境中指定，必须注意确保为速度正确编译，否则会观察到明显的减速。
单程解决方案需要额外的代码来处理未使用的（但最初分配的）内存。
NumPy 方法不是 return 输入的视图，而是一个副本，作为 advanced indexing:

arr = np.arange(100)
k = 50
print('`arr[arr > k]` is a copy: ', arr[arr > k].base is None)
# `arr[arr > k]` is a copy:  True
print('`arr[np.where(arr > k)]` is a copy: ', arr[np.where(arr > k)].base is None)
# `arr[np.where(arr > k)]` is a copy:  True
print('`arr[:k]` is a copy: ', arr[:k].base is None)
# `arr[:k]` is a copy:  False

(已编辑：包含基于 np.nonzero() 的解决方案和修复内存泄漏并避免在单遍 Cython/Numba 版本中复制，包含两遍 Cython/Numba 版本——基于@ShadowRanger、@PaulPanzer、@max9111 和@DavidW 评论。）

过滤 NumPy 数组

Filtering a NumPy Array

python

numpy

cython

numba

总结

定义

计时基准

作为输入大小的函数

作为填充函数

内存注意事项

备注