Numba guvectorize 与 njit：处理不同尺寸数组返回的策略

引言：Numba guvectorize 与变长数组返回的挑战

在使用 Numba 对 Python 函数进行性能优化时，guvectorize 装饰器是一个强大的工具，它允许我们创建广义的向量化函数（Generalized Universal Functions, GUFuncs），从而在 Numba JIT 编译的代码中实现数组操作的并行化。然而，当函数的输出数组尺寸与输入数组尺寸不同时，尤其是在尝试直接返回一个固定尺寸（例如，用于统计唯一值出现次数的 257 长度数组）的新数组时，开发者可能会遇到编译错误或行为不符合预期的问题。这通常是由于对 guvectorize 的设计原理和其签名规范理解不足所致。

理解 guvectorize 的设计哲学与局限

guvectorize 的核心思想是为 NumPy 的 ufunc 机制提供一个广义的扩展。它旨在处理具有“核心”维度（core dimensions）的数组操作，这些核心维度在函数内部被处理，而其他“批次”维度（batch dimensions）则由 Numba 自动进行循环和并行化。其签名字符串 "(n) -> (m)" 定义了输入和输出的核心维度，其中 n 和 m 代表核心维度的长度。

然而，guvectorize 的一个关键限制是它并不支持直接返回一个形状与输入核心维度无关的新数组。具体来说：

1. 输出数组形状的推导： guvectorize 期望输出数组的形状能够根据输入数组的形状和签名字符串推导出来。它不是为了返回一个完全独立、固定尺寸的数组而设计的。
2. void 返回类型： guvectorize 函数通常应声明为 void 返回类型。这意味着函数内部不应使用 return 语句显式返回任何值。相反，输出数组应该作为函数的参数传入，并在函数内部进行修改（in-place modification）。
3. 并行化模型： guvectorize 的并行化是基于批次维度进行的。对于简单的 1D 数组处理，如果不存在批次维度需要并行，其优势可能不如 njit 明显。

因此，尝试在 guvectorize 函数内部创建并返回一个新数组（如 count = np.zeros(...) 并 return count）是错误的用法，会导致编译失败或运行时异常。

guvectorize 的正确实践：通过参数传递输出数组

要正确使用 guvectorize 来实现类似统计唯一值的功能，同时返回一个固定尺寸的数组，正确的做法是预先分配好输出数组，并将其作为参数传递给 guvectorize 函数。函数内部将直接修改这个传入的数组。

以下是实现字节数组中唯一值计数并返回固定长度计数数组的正确 guvectorize 示例：

import numpy as np
import numba as nb

@nb.guvectorize("void(uint8[:], uint64[:])", "(n),(m)", target="cpu")
def count_occurrences(byte_view, count):
    """
    Counts the occurrences of each element in a byte array and updates the count array in-place.

    Parameters:
    byte_view (np.uint8[:]): The input byte array.
    count (np.uint64[:]): The output array to store counts. It should be pre-allocated.
                           The first element (index 0) is typically unused for convenience
                           when counting values from 0-255.
    """
    # Ensure the count array is initialized to zeros if not already.
    # For guvectorize, it's generally assumed the caller handles initialization.
    # If not, a loop to zero it out might be needed, but often unnecessary
    # if the array is freshly created with np.zeros.

    # Iterate over each byte in the input view and increment the corresponding count.
    # We add 1 to the byte value to account for the leading zero in the count array.
    for idx in byte_view: 
        count[1 + idx] += 1

# Example usage:
sample = np.random.randint(1, 100, 100, dtype=np.uint8)

# Pre-allocate the output array.
# It has a length of 257 (1 for index 0, and 256 for values 0-255).
counts = np.zeros(1 + 256, dtype=np.uint64)

# Call the guvectorized function. The 'counts' array is modified in-place.
count_occurrences(sample, counts)

print("Sample input:", sample[:10])
print("Counts output:", counts[1:10]) # Display counts for values 0-9
print("Total elements counted:", np.sum(counts[1:])) # Should match sample.size

代码解析：

• @nb.guvectorize("void(uint8[:], uint64[:])", "(n),(m)", target="cpu"):
  • 第一个参数 void(uint8[:], uint64[:]) 定义了函数的类型签名。void 表示函数不返回任何值。uint8[:] 和 uint64[:] 分别表示第一个输入参数是 uint8 类型的一维数组，第二个参数是 uint64 类型的一维数组。第二个参数 count 在这里被视为输出参数。
  • 第二个参数 "(n),(m)" 定义了核心维度。(n) 表示第一个输入数组 byte_view 有一个核心维度 n。(m) 表示第二个数组 count 也有一个核心维度 m。Numba 会确保在调用时，count 数组的尺寸与此 m 维度匹配。
  • target="cpu" 指定在 CPU 上执行。
• 函数体： 函数体内部的逻辑与原始意图相同，遍历 byte_view 中的每个元素，并更新 count 数组中对应位置的值。
• 调用方式： 在调用 count_occurrences 之前，必须先使用 np.zeros 等方法预先分配好 counts 数组，并将其作为参数传入。函数执行后，counts 数组的内容会被更新。

guvectorize 与 njit 的选择考量

虽然上述 guvectorize 的实现是正确的，但对于这种特定的任务（简单的 1D 数组统计，且输出数组形状固定），guvectorize 的优势可能并不突出。实际上，对于许多需要返回新数组且形状不直接依赖于 guvectorize 核心维度推导的场景，@nb.njit 装饰器可能是一个更简单、更直观的选择。

即构数智人

即构数智人是由即构科技推出的AI虚拟数字人视频创作平台，支持数字人形象定制、短视频创作、数字人直播等。

36

查看详情

@nb.njit 允许函数直接创建并返回一个新的 NumPy 数组，而无需考虑 guvectorize 的复杂签名和 void 返回限制。

import numpy as np
import numba as nb

@nb.njit
def count_occurrences_njit(byte_view):
    """
    Counts the occurrences of each element in a byte array and returns a new array with the counts.
    This version uses njit, allowing direct return of a new array.
    """
    # Create and initialize the count array directly within the njit function
    count = np.zeros(1 + 256, dtype=np.uint64)
    for idx in byte_view:
        count[1 + idx] += 1
    return count

# Example usage with njit:
sample_njit = np.random.randint(1, 100, 100, dtype=np.uint8)
counts_njit = count_occurrences_njit(sample_njit)

print("\nSample input (njit):", sample_njit[:10])
print("Counts output (njit):", counts_njit[1:10])
print("Total elements counted (njit):", np.sum(counts_njit[1:]))

何时选择 guvectorize：

• 当你需要创建广义的 ufunc，并且你的操作可以被分解为独立的核心维度操作，Numba 可以通过批次维度进行并行化时。
• 当输出数组的形状可以根据输入数组的形状和核心维度签名进行推导时。
• 当你希望将函数的计算结果直接写入预分配的输出数组中，以避免内存分配开销时。

何时选择 njit：

• 当你的函数逻辑相对简单，不需要 guvectorize 提供的复杂批次并行化机制时。
• 当你需要直接从函数中返回新创建的、形状灵活或固定但与输入核心维度无关的数组时。
• 当你只需要对 Python 代码进行即时编译以提高性能时。

对于本教程中的计数问题，由于其不涉及复杂的批次维度并行化，且输出数组形状固定，njit 的实现可能更为简洁和直观。

总结与最佳实践

在使用 Numba 优化代码时，理解不同装饰器的设计目的和适用场景至关重要：

1. guvectorize 的核心用途：它主要用于创建广义的 ufunc，实现对数组核心维度的操作，并利用 Numba 的并行化机制处理批次维度。它的函数签名严格，且通常要求函数返回 void，通过参数传递并修改输出数组。
2. 处理不同尺寸输出：如果 guvectorize 函数需要产生一个与输入尺寸不同的数组，正确的做法是预先分配该输出数组，并将其作为参数传入函数进行修改。
3. njit 的灵活性：对于许多场景，特别是当函数需要直接创建并返回一个新数组，且其形状不严格依赖于 guvectorize 的核心维度推导时，@nb.njit 是一个更简单、更灵活的选择。
4. 选择合适的工具：根据你的具体需求——是需要广义的 ufunc 和批次并行化，还是仅仅需要编译一个高性能的 Python 函数并灵活处理返回值——来选择 guvectorize 或 njit。

通过深入理解这些 Numba 装饰器的特性，开发者可以更有效地编写高性能的 Python 代码。

以上就是Numba guvectorize 与 njit：处理不同尺寸数组返回的策略的详细内容，更多请关注php中文网其它相关文章！