筛选NumPy数组

2022-04-04 00:00:00 python numpy cython numba

问题描述

假设我有一个NumPy数组arr，我想要按元素进行筛选，例如 我只想获取低于某个阈值k的值。

有几种方法，例如：

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

哪一个是最快的？内存效率如何？

(编辑：根据@ShadowRanger评论增加np.nonzero())

解决方案

摘要

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

这里测试的不完全是一般过滤，而只是应用阈值，这具有计算条件相当快的显著特征。如果该条件意味着昂贵的计算，则将获得非常不同的结果。

基本上，两遍加速(无论是使用Numba还是Cython--只要您事先知道类型)将是最快和更高效的内存，除了非常大的输入，对于这种情况，单遍Numba/Cython更快(代价是使用更大的临时内存)。使用np.where()/np.nonzero()而不是直接使用掩码可能会导致计算略微加快，而且只要输出的大小小于输入的50%，通常不会有任何影响(可能除了较大的临时内存占用)。np.fromiter()方法要慢得多，但不会生成大型临时对象。

定义

  1. 使用生成器：
def filter_fromiter(arr, k):
    return np.fromiter((x for x in arr if x < k), dtype=arr.dtype)
  1. 使用布尔掩码切片：
def filter_mask(arr, k):
    return arr[arr < k]
  1. 使用np.where()：
def filter_where(arr, k):
    return arr[np.where(arr < k)]
  1. 使用np.nonzero()
def filter_nonzero(arr, k):
    return arr[np.nonzero(arr < k)]
  1. 使用基于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'
  1. 使用基于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的解决方案相比，当填充小于60%时，np.where()/np.nonzero()解决方案优于布尔掩码切片，之后布尔掩码切片速度最快。

(完整代码可用here)

内存注意事项

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

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

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

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

备注

  • 在指定不同的过滤条件时，生成器方法也是最灵活的
  • Cython解决方案需要指定数据类型以使其速度更快，否则需要为多个类型分派额外的工作
  • 对于Numba和Cython，可以将筛选条件指定为泛型函数(因此不需要进行硬编码)，但必须在其各自的环境中指定，并且必须注意确保正确编译以提高速度，否则会出现明显的减速。
  • 单遍解决方案需要额外的代码来处理未使用的(但最初分配的)内存。
  • NumPy方法不返回输入的视图，而是返回一个副本，结果是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的评论。)

相关文章