Skip to content

1.11。使用@stencil装饰器

原文： http://numba.pydata.org/numba-doc/latest/user/stencil.html

模板是一种常见的计算模式，其中数组元素根据称为模板内核的某种固定模式进行更新。 Numba 提供@stencil装饰器，以便用户可以轻松指定模板内核，然后 Numba 生成将该内核应用于某个输入数组所需的循环代码。因此，模板装饰器允许更清晰，更简洁的代码，并且与一起，并行 jit 选项通过模板执行的并行化实现更高的性能。

1.11.1。基本用法

@stencil装饰器的示例用法：

from numba import stencil

@stencil
def kernel1(a):
    return 0.25 * (a[0, 1] + a[1, 0] + a[0, -1] + a[-1, 0])

模板内核由看起来像标准 Python 函数定义的内容指定，但是在数组索引方面有不同的语义。模板生成与输入数组具有相同大小和形状的输出数组，尽管取决于内核定义可能具有不同的类型。从概念上讲，模板内核对输出数组中的每个元素运行一次。模板内核的返回值是写入该特定元素的输出数组的值。

参数a表示应用内核的输入数组。索引到此数组是针对正在处理的输出数组的当前元素进行的。例如，如果正在处理元素(x, y)，则模板内核中的a[0, 0]对应于输入数组中的a[x + 0, y + 0]。类似地，模板内核中的a[-1, 1]对应于输入数组中的a[x - 1, y + 1]。

根据指定的内核，内核可能不适用于输出数组的边界，因为这可能导致输入数组被越界访问。模板装饰器处理这种情况的方式取决于选择 func_or_mode 。默认模式是模板修饰器将输出数组的边框元素设置为零。

要在输入数组上调用模板，请将模板调用为常规函数，并将输入数组作为参数传递。例如，使用上面定义的内核：

>>> import numpy as np
>>> input_arr = np.arange(100).reshape((10, 10))
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14, 15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35, 36, 37, 38, 39],
       [40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
       [50, 51, 52, 53, 54, 55, 56, 57, 58, 59],
       [60, 61, 62, 63, 64, 65, 66, 67, 68, 69],
       [70, 71, 72, 73, 74, 75, 76, 77, 78, 79],
       [80, 81, 82, 83, 84, 85, 86, 87, 88, 89],
       [90, 91, 92, 93, 94, 95, 96, 97, 98, 99]])
>>> output_arr = kernel1(input_arr)
array([[  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.],
       [  0.,  11.,  12.,  13.,  14.,  15.,  16.,  17.,  18.,   0.],
       [  0.,  21.,  22.,  23.,  24.,  25.,  26.,  27.,  28.,   0.],
       [  0.,  31.,  32.,  33.,  34.,  35.,  36.,  37.,  38.,   0.],
       [  0.,  41.,  42.,  43.,  44.,  45.,  46.,  47.,  48.,   0.],
       [  0.,  51.,  52.,  53.,  54.,  55.,  56.,  57.,  58.,   0.],
       [  0.,  61.,  62.,  63.,  64.,  65.,  66.,  67.,  68.,   0.],
       [  0.,  71.,  72.,  73.,  74.,  75.,  76.,  77.,  78.,   0.],
       [  0.,  81.,  82.,  83.,  84.,  85.,  86.,  87.,  88.,   0.],
       [  0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.,   0.]])
>>> input_arr.dtype
dtype('int64')
>>> output_arr.dtype
dtype('float64')

请注意，模板装饰器已确定指定的模板内核的输出类型为float64，因此将输出数组创建为float64，而输入数组的类型为int64。

1.11.2。模板参数

模板内核定义可以采用以下规定的任意数量的参数。第一个参数必须是数组。输出数组的大小和形状将与第一个参数的大小和形状相同。其他参数可以是标量或数组。对于数组参数，这些数组必须至少与每个维度中的第一个参数（数组）一样大。所有此类输入数组参数的数组索引都是相对的。

1.11.3。内核形状推理和边界处理

在上面的示例中，在大多数情况下，模板内核中的数组索引将专门使用Integer文字。在这种情况下，模板装饰器能够分析模板内核以确定其大小。在上面的例子中，模板装饰器确定内核是3 x 3的形状，因为索引-1到1用于第一和第二维度。请注意，模板装饰器也可以正确处理非对称和非方形模板内核。

根据模板内核的大小，模板装饰器能够计算输出数组中边框的大小。如果将内核应用于输入数组的某个元素会导致索引超出范围，那么该元素属于输出数组的边界。在上面的示例中，在每个维度中访问点-1和+1，因此输出数组在所有维度中具有大小为 1 的边界。

如果可能，并行模式能够从简单表达式推断内核索引作为常量。例如：

@njit(parallel=True)
def stencil_test(A):
    c = 2
    B = stencil(
        lambda a, c: 0.3 * (a[-c+1] + a[0] + a[c-1]))(A, c)
    return B

1.11.4。模板装饰器选项

注意

模板装饰器可以在将来增强，以提供用于边界处理的附加机制。目前，只实施了一种行为，"constant"（详见下文func_or_mode）。

1.11.4.1。 neighborhood

有时用Integer文字专门编写模板内核可能不方便。例如，假设我们想计算一组时间序列数据的 30 天移动平均值。可以编写(a[-29] + a[-28] + ... + a[-1] + a[0]) / 30，但模板装饰器使用neighborhood选项提供更简洁的形式：

@stencil(neighborhood = ((-29, 0),))
def kernel2(a):
    cumul = 0
    for i in range(-29, 1):
        cumul += a[i]
    return cumul / 30

邻域选项是元组的元组。外元组的长度等于输入数组的维数。内元组的长度始终为 2，因为外元组的每个元素对应于相应维度中使用的最小和最大索引偏移量。

如果用户指定了邻域但内核访问了指定邻域之外的元素，则行为未定义。

1.11.4.2。 func_or_mode

可选的func_or_mode参数控制如何处理输出数组的边框。目前，只有一个受支持的值，"constant"。在constant模式下，在内核访问输入数组有效范围之外的元素的情况下，不应用模板内核。在这种情况下，输出数组中的那些元素被赋值为常量值，由cval参数指定。

1.11.4.3。 cval

可选的 cval 参数默认为零，但可以设置为任何所需的值，如果func_or_mode参数设置为constant，则该值用于输出数组的边界。在所有其他模式中忽略 cval 参数。 cval 参数的类型必须与模板内核的返回类型匹配。如果用户希望输出数组是从特定类型构造的，那么他们应该确保模板内核返回该类型。

1.11.4.4。 standard_indexing

默认情况下，模板内核中的所有数组访问都作为相对索引处理，如上所述。然而，有时将辅助数组（例如权重数组）传递给模板内核并使该数组使用标准 Python 索引而不是相对索引可能是有利的。为此，有一个模板装饰器选项standard_indexing，其值是一个字符串的集合，其名称与模板函数的参数相匹配，这些参数将使用标准 Python 索引而不是相对索引来访问：

@stencil(standard_indexing=("b",))
def kernel3(a, b):
    return a[-1] * b[0] + a[0] + b[1]

1.11.5。 StencilFunc

模板装饰器返回StencilFunc类型的可调用对象。 StencilFunc对象包含许多属性，但用户可能感兴趣的唯一属性是neighborhood属性。如果将neighborhood选项传递给模板装饰器，则提供的邻域将存储在此属性中。否则，在首次执行或编译时，系统如上所述计算邻域，然后将计算的邻域存储到该属性中。然后，如果用户希望验证所计算的邻域是正确的，则用户可以检查该属性。

1.11.6。模板调用选项

在内部，模板装饰器将指定的模板内核转换为常规的 Python 函数。此函数将具有与模板内核定义中指定的参数相同的参数，但也将包含以下可选参数。

1.11.6.1。 out

可选的out参数被添加到 Numba 生成的每个模板函数中。如果指定，out参数告诉 Numba 用户正在提供他们自己的预分配数组以用于模板的输出。在这种情况下，模板函数不会分配自己的输出数组。用户应确保模板内核的返回类型可以安全地转换为 Numpy ufunc 强制转换规则之后的用户指定输出数组的元素类型。

示例用法如下所示：

>>> import numpy as np
>>> input_arr = np.arange(100).reshape((10, 10))
>>> output_arr = np.full(input_arr.shape, 0.0)
>>> kernel1(input_arr, out=output_arr)
回到顶部