在这篇文章中，我们将探讨四种数组计算技术，以及它们如何协同工作以取得强大的结果。

1. Numba 的 stencil 装饰器用于创建局部计算内核
2. Numba 的即时 (JIT) 编译器用于 Python 中的数组计算
3. Dask 数组用于跨多个块并行化数组计算
4. NumPy 的广义通用函数 (gufuncs) 将所有内容连接起来

最后，我们将展示新手开发者如何编写少量 Python 代码，以高效地对大量数据执行局部计算。特别是，我们将编写一个简单的函数来平滑图像，并将其并行应用于大量图像堆栈。

以下是完整的代码，我们将在下面逐部分深入探讨。

import numba

@numba.stencil
def _smooth(x)
return (x[-1, -1] + x[-1, 0] + x[-1, 1] +
x[ 0, -1] + x[ 0, 0] + x[ 0, 1] +
x[ 1, -1] + x[ 1, 0] + x[ 1, 1]) // 9


@numba.guvectorize(
[(numba.int8[:, :], numba.int8[:, :])],
'(n, m) -> (n, m)'
)
def smooth(x, out)
out[:] = _smooth(x)


# If you want fake data
import dask.array as da
x = da.ones((1000000, 1000, 1000), chunks=('auto', -1, -1), dtype='int8')

# If you have actual data
import dask_image
x = dask_image.imread.imread('/path/to/*.png')

y = smooth(x)
# dask.array<transpose, shape=(1000000, 1000, 1000), dtype=int8, chunksize=(125, 1000, 1000)>

注意：上面的 smooth 函数在图像处理领域通常被称为二维均值滤波器。

现在，让我们对此进行一些分解

Numba Stencils

文档：: https://numba.pydata.org/numba-doc/dev/user/stencil.html

许多数组计算函数只在数组的局部区域上操作。这在图像处理、信号处理、仿真、微分方程求解、异常检测、时间序列分析等方面很常见。通常我们编写的代码如下所示

def _smooth(x)
out = np.empty_like(x)
for i in range(1, x.shape[0] - 1)
for j in range(1, x.shape[1] - 1)
out[i, j] = x[i + -1, j + -1] + x[i + -1, j + 0] + x[i + -1, j + 1] +
x[i + 0, j + -1] + x[i + 0, j + 0] + x[i + 0, j + 1] +
x[i + 1, j + -1] + x[i + 1, j + 0] + x[i + 1, j + 1]) // 9

return out

或者类似的代码。`numba.stencil` 装饰器使编写起来更容易一些。您只需写下每个元素上发生的情况，剩下的交给 Numba 处理。

@numba.stencil
def _smooth(x)
return (x[-1, -1] + x[-1, 0] + x[-1, 1] +
x[ 0, -1] + x[ 0, 0] + x[ 0, 1] +
x[ 1, -1] + x[ 1, 0] + x[ 1, 1]) // 9

Numba JIT

文档： http://numba.pydata.org/

当我们在 NumPy 数组上运行此函数时，我们发现它很慢，运行速度与 Python 速度相当。

x = np.ones((100, 100))
timeit _smooth(x)
527 ms ± 44.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

但是如果使用 Numba 对此函数进行 JIT 编译，那么它的运行速度会更快。

@numba.njit
def smooth(x)
return _smooth(x)

%timeit smooth(x)
70.8 µs ± 6.38 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

对于那些计算的人来说，这快了 1000 多倍！

注意：此函数已存在于 `scipy.ndimage.uniform_filter` 中，其运行速度相同。

Dask 数组

文档： https://docs.dask.org.cn/en/latest/array.html

在这些应用中，人们通常有很多这样的数组，并且他们想将此函数应用于所有数组。原则上，他们可以使用 for 循环来完成此操作。

from glob import glob
import skimage.io

for fn in glob('/path/to/*.png')
img = skimage.io.imread(fn)
out = smooth(img)
skimage.io.imsave(fn.replace('.png', '.out.png'), out)

如果他们想并行执行此操作，他们可能会使用 `multiprocessing` 或 `concurrent.futures` 模块。如果他们想在集群上执行此操作，他们可以使用 PySpark 或其他系统重写代码。

或者，他们可以使用 Dask 数组，它将处理管道化和并行化（单机或集群），同时看起来仍然大部分像 NumPy 数组。

import dask_image
x = dask_image.imread('/path/to/*.png') # a large lazy array of all of our images
y = x.map_blocks(smooth, dtype='int8')

然后，由于 Dask 数组的每个块都只是 NumPy 数组，我们可以使用 `map_blocks` 函数将此函数应用于所有图像，然后将其保存。

这很好，但让我们进一步深入，讨论 NumPy 的广义通用函数。

广义通用函数

Numba 文档： https://numba.pydata.org/numba-doc/dev/user/vectorize.html

NumPy 文档： https://docs.scipy.org.cn/doc/numpy-1.16.0/reference/c-api.generalized-ufuncs.html

广义通用函数 (gufunc) 是一个普通函数，已使用类型和维度信息进行了标注。例如，我们可以如下将我们的 smooth 函数重新定义为 gufunc

@numba.guvectorize(
[(numba.int8[:, :], numba.int8[:, :])],
'(n, m) -> (n, m)'
)
def smooth(x, out)
out[:] = _smooth(x)

此函数知道它消耗一个 int8 的二维数组，并生成一个具有相同维度的 int8 的二维数组。

这种标注是一个小小的改动，但它为 Dask 等其他系统提供了足够的信息来智能地使用它。我们不必调用像 `map_blocks` 这样的函数，而是可以直接使用该函数，就像我们的 Dask 数组只是一个非常大的 NumPy 数组一样。

# Before gufuncs
y = x.map_blocks(smooth, dtype='int8')

# After gufuncs
y = smooth(x)

这很好。如果您使用 gufunc 语义编写库代码，那么该代码就可以直接与 Dask 等系统一起工作，而无需构建明确的并行计算支持。这大大简化了用户的工作。

最终结果

让我们再看一次完整的示例。

import numba
import dask.array as da

@numba.stencil
def _smooth(x)
return (x[-1, -1] + x[-1, 0] + x[-1, 1] +
x[ 0, -1] + x[ 0, 0] + x[ 0, 1] +
x[ 1, -1] + x[ 1, 0] + x[ 1, 1]) // 9


@numba.guvectorize(
[(numba.int8[:, :], numba.int8[:, :])],
'(n, m) -> (n, m)'
)
def smooth(x, out)
out[:] = _smooth(x)

x = da.ones((1000000, 1000, 1000), chunks=('auto', -1, -1), dtype='int8')
smooth(x)

这段代码对于新手用户来说相当易于理解。他们可能不理解 gufunc、Dask 数组或 JIT 编译的内部细节，但他们很可能会复制、粘贴和修改上面的示例以满足他们的需求。

他们确实想改变的部分很容易改变，比如 stencil 计算，以及创建自己的数据数组。

这个工作流程高效且可扩展，使用了底层编译代码以及可能包含数千台计算机的集群。

可以改进的地方

有一些地方可以让这个工作流程更好。

1. 如果不必在 guvectorize 中指定 dtypes，而是根据传入的类型进行专门化，那就太好了。numba/numba #2979
2. 使用numba.cuda.jit支持 GPU 加速 stencil 计算。Stencil 计算是 GPU 加速的明显候选，这是一个很好的可访问点，新手用户可以以一种对自动化系统来说足够受限的方式指定他们想要的东西，以便相对容易地将其重写为 CUDA。numba/numba 3915
3. 如果能够像这样直接在 stencil 函数顶部应用 @guvectorize 装饰器，那就更好了。
4. @numba.guvectorize(...)
   @numba.stencil
   def average(x)
   ...
   
7. 而不是有两个函数。numba/numba #3914
9. 您可能已经注意到，我们的 guvectorize 函数必须将其结果分配给一个 out 参数。
10. @numba.guvectorize(
    [(numba.int8[:, :], numba.int8[:, :])],
    '(n, m) -> (n, m)'
    )
    def smooth(x, out)
    out[:] = _smooth(x)
    
13. 也许直接返回输出会更好
15. def smooth(x)
    return _smooth(x)
    
18. numba/numba #3916
20. dask-image 库可以使用一个 imsave 函数
21. dask/dask-image #110

理想结果

有了这些，我们就可以将上面的代码写成如下所示

# This is aspirational

import numba
import dask_image

@numba.guvectorize(
[(numba.int8[:, :], numba.int8[:, :])],
signature='(n, m) -> (n, m)',
target='gpu'
)
@numba.stencil
def smooth(x)
return (x[-1, -1] + x[-1, 0] + x[-1, 1] +
x[ 0, -1] + x[ 0, 0] + x[ 0, 1] +
x[ 1, -1] + x[ 1, 0] + x[ 1, 1]) // 9

x = dask_image.io.imread('/path/to/*.png')
y = smooth(x)
dask_image.io.imsave(y, '/path/to/out/*.png')

更新：现在支持 GPU！

写完这篇博文后，我做了一个小更新，使用numba.cuda.jit在 GPU 上实现了相同的平滑函数，代码复杂度仅略有增加，但速度提高了 200 倍。

该 notebook 在这里.