使用 Numba 进行模板计算
内容
实时笔记本
你可以在 live session 
中运行此笔记本，或查看 Github 上的内容。
使用 Numba 进行模板计算¶
本笔记本结合了 Numba ，一个高性能的 Python 编译器，与 Dask 数组。
特别是，我们展示了两个Numba特性，以及它们如何与Dask结合：
Numba 的 stencil decorator
NumPy 的 通用函数
这篇文章最初作为博客文章发布 在这里
Numba Stencils 简介¶
许多数组计算函数仅对数组的局部区域进行操作。这在图像处理、信号处理、模拟、微分方程求解、异常检测、时间序列分析等领域很常见。通常我们会编写如下代码：
[ ]:
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 会处理其余部分。
[ ]:
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
当我们在一个NumPy数组上运行这个函数时，我们发现它运行得很慢，操作速度接近Python的速度。
[ ]:
import numpy as np
x = np.ones((100, 100))
%timeit _smooth(x)
但如果我们使用 Numba 即时编译这个函数，那么它的运行速度会更快。
[ ]:
@numba.njit
def smooth(x):
return _smooth(x)
%timeit smooth(x)
对于那些计算的人来说，这比原来快了1000倍以上！
注意：此功能已作为 ``scipy.ndimage.uniform_filter`` 存在，其操作速度相同。
Dask 数组¶
在这些应用中，人们通常有许多这样的数组，他们希望将此函数应用于所有这些数组。原则上，他们可以使用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中的广义通用函数。
因为我们附近没有一堆图片，我们将制作一个具有类似结构的随机数组。
[ ]:
import dask.array as da
x = da.random.randint(0, 127, size=(10000, 1000, 1000), chunks=('64 MB', None, None), dtype='int8')
x
广义通用函数¶
Numba 文档: https://numba.pydata.org/numba-doc/dev/user/vectorize.html
NumPy 文档: https://numpy.org/doc/stable/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)
y
这很不错。如果你用 gufunc 语义编写库代码，那么这些代码就可以直接与 Dask 这样的系统一起工作，而无需你专门为并行计算构建显式支持。这使得用户的生活变得更加轻松。
启动 Dask 客户端以使用仪表板¶
启动 Dask 客户端是可选的。它将启动仪表板，这对于深入了解计算非常有用。
[ ]:
from dask.distributed import Client, progress
client = Client(threads_per_worker=4,
n_workers=1,
processes=False,
memory_limit='4GB')
client
[ ]:
y.max().compute()
GPU 版本¶
Numba 还支持在兼容的 GPU 设备上使用 CUDA 进行即时编译。
这比在单个 V100 GPU 上使用 numba.cuda.jit 的 CPU 快了大约 200 倍。
import numba.cuda
@numba.cuda.jit
def smooth_gpu(x, out):
i, j = cuda.grid(2)
n, m = x.shape
if 1 <= i < n - 1 and 1 <= j < m - 1:
out[i, j] = (x[i - 1, j - 1] + x[i - 1, j] + x[i - 1, j + 1] +
x[i , j - 1] + x[i , j] + x[i , j + 1] +
x[i + 1, j - 1] + x[i + 1, j] + x[i + 1, j + 1]) // 9
import cupy, math
x_gpu = cupy.ones((10000, 10000), dtype='int8')
out_gpu = cupy.zeros((10000, 10000), dtype='int8')
# I copied the four lines below from the Numba docs
threadsperblock = (16, 16)
blockspergrid_x = math.ceil(x_gpu.shape[0] / threadsperblock[0])
blockspergrid_y = math.ceil(x_gpu.shape[1] / threadsperblock[1])
blockspergrid = (blockspergrid_x, blockspergrid_y)
smooth_gpu[blockspergrid, threadsperblock](x_gpu, out_gpu)
完整笔记本 在此