Dask 工作笔记

将 Dask Array 与 Numba Stencils 结合使用

作者：Matthew Rocklin

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

  1. Numba 的 stencil 装饰器，用于创建局部计算核
  2. Numba 的即时 (JIT) 编译器，用于在 Python 中进行数组计算
  3. Dask Array，用于在多个块上并行化数组计算
  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 Array

文档： 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 Array，它将处理流水线和并行化（单机或在集群上），同时看起来仍然像一个 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 Array 的每个块都只是 NumPy 数组，我们可以使用 map_blocks 函数将这个函数应用到所有图像上，然后保存它们。

这已经不错了，但让我们更进一步，讨论 NumPy 的通用通用函数。

通用通用函数 (Generalized Universal Functions)

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)

这个函数知道它接收一个 2d 的 int8 数组，并生成一个具有相同维度的 2d 的 int8 数组。

这种类型的注解只是一个小的改动，但它给了 Dask 等其他系统足够的信息来智能地使用它。我们不必调用 map_blocks 之类的函数，而是可以直接使用该函数，就像我们的 Dask Array 是一个非常大的 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)

这段代码对于新手用户来说很容易上手。他们可能不理解 gufuncs 或 Dask arrays 或 JIT 编译的内部细节，但他们很可能能够复制-粘贴-修改上面的示例来满足他们的需求。

他们想要更改的部分很容易更改，比如模板计算和创建他们自己的数据数组。

这个工作流程高效且可扩展，使用低级编译代码，并可能利用数千台计算机的集群。

哪些可以做得更好

有几件事可以让这个工作流程更顺畅。

  1. 最好不必在 guvectorize 中指定 dtypes，而是根据传入的类型进行特化。numba/numba #2979

  2. 支持使用 numba.cuda.jit 对模板计算进行 GPU 加速。模板计算是 GPU 加速的明显候选者，这是一个很好的、易于入门的点，新手用户可以用一种足够受约束的方式指定他们想要的东西，以便自动化系统可以相对容易地将其重写为 CUDA。numba/numba 3915

  3. 如果能够直接将 @guvectorize 装饰器应用到模板函数之上，像这样，那就更好了。

    @numba.guvectorize(...)
@numba.stencil
def average(x):
    ...

    而不是有两个函数。numba/numba #3914

  4. 您可能已经注意到，我们的 guvectorize 函数必须将其结果赋值给一个 out 参数。

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

    或许，直接返回输出会更好

    def smooth(x):
    return _smooth(x)

    numba/numba #3916

  5. dask-image 库可以使用一个 imsave 函数

    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 上实现了同样的 smooth 函数，仅增加了少许代码复杂性就实现了 200 倍的加速。

该笔记本在此.

由 Disqus 提供支持的博客评论