使用 Numba 加速数组统计：guvectorize 的正确使用姿势-Python教程-PHP中文网

使用 Numba 加速数组统计：guvectorize 的正确使用姿势

花韻仙語

发布： 2025-10-07 10:46:28

原创

141人浏览过

使用 numba 加速数组统计：guvectorize 的正确使用姿势

第一段引用上面的摘要：

本文旨在阐述如何使用 Numba 的 guvectorize 装饰器来加速数组统计计算，特别是当输出数组的形状与输入数组不同时。我们将通过示例代码详细解释 guvectorize 的正确用法，并讨论其与 njit 的区别与适用场景，帮助读者理解并掌握 Numba 优化数组操作的技巧。

理解 guvectorize 的限制与正确用法

numba.guvectorize 是一个强大的工具，用于将标量函数推广到可以在 NumPy 数组上操作的通用函数。然而，它也存在一些限制。其中一个关键限制是，使用 guvectorize 时，无法直接返回形状与输入数组不同的数组。这意味着输出数组的形状必须预先确定，并且通常需要作为输入参数传递给函数。

根本原因在于 guvectorize 的设计目标是并行化操作，而返回动态形状的数组会使得并行化变得复杂。因此，guvectorize 的函数应该修改传入的数组，而不是返回新的数组。

示例：统计字节数组中元素的出现次数

假设我们需要编写一个函数，用于统计一个字节数组中每个元素的出现次数，并将结果存储在一个长度为 257 的数组中（包括 0 到 256 的计数）。以下是如何使用 guvectorize 实现此功能的示例：

import numpy as np
import numba as nb

@nb.guvectorize("void(uint8[:], uint64[:])", "(n),(m)", target="cpu")
def count_occurrences(byte_view, count):
    """
    统计字节数组中每个元素的出现次数，并将结果存储在 count 数组中。
    """
    for idx in byte_view:
        count[1 + idx] += 1

# 示例用法
sample = np.random.randint(1, 100, 100, dtype=np.uint8)

# 初始化计数数组
counts = np.zeros(1 + 256, dtype=np.uint64)

# 调用 guvectorize 函数
count_occurrences(sample, counts)

print(counts)

登录后复制

代码解释：

@nb.guvectorize("void(uint8[:], uint64[:])", "(n),(m)", target="cpu"): 这是 guvectorize 装饰器。
- "void(uint8[:], uint64[:])"：指定了函数签名。它表示该函数接受一个 uint8 类型的 1D 数组和一个 uint64 类型的 1D 数组作为输入，并且不返回任何值（void）。
- "(n),(m)"：指定了输入和输出数组的布局。 (n) 表示第一个输入数组的形状是 (n)，(m) 表示第二个输入数组的形状是 (m)。
- target="cpu"：指定了目标平台为 CPU。
count_occurrences(byte_view, count): 这是被 guvectorize 装饰的函数。

SpeakingPass-打造你的专属雅思口语语料
使用chatGPT帮你快速备考雅思口语，提升分数

25

查看详情
- byte_view：输入的字节数组。
- count：用于存储元素出现次数的数组。注意，count 数组是作为输入参数传递的，并且在函数内部被修改。
- 函数内部循环遍历 byte_view 中的每个元素，并更新 count 数组中对应元素的计数。count[1 + idx] += 1 的意思是将 byte_view 中值为 idx 的元素的个数加一。因为 count 的第一个元素（索引为 0）未被使用，所以索引需要加 1。
counts = np.zeros(1 + 256, dtype=np.uint64): 在调用 guvectorize 函数之前，我们需要初始化 count 数组。
count_occurrences(sample, counts): 调用 guvectorize 函数，将 sample 数组和 counts 数组作为输入传递。

guvectorize vs. njit：选择合适的工具

虽然上面的示例展示了如何使用 guvectorize 来解决问题，但值得注意的是，在这种情况下，使用 numba.njit 可能更简单，并且性能差异可能很小。

import numpy as np
import numba as nb

@nb.njit
def count_occurrences_njit(byte_view):
    """
    使用 njit 统计字节数组中每个元素的出现次数，并返回结果数组。
    """
    count = np.zeros(1 + 256, dtype=np.uint64)
    for idx in byte_view:
        count[1 + idx] += 1
    return count

# 示例用法
sample = np.random.randint(1, 100, 100, dtype=np.uint8)

# 调用 njit 函数
counts = count_occurrences_njit(sample)

print(counts)

登录后复制

选择建议：