dask-使用Numba进行模版计算

使用Numba进行模具计算

Numba徽标

这款笔记本结合Numba,高性能的Python编译器,用DASK阵列。

特别是,我们展示了Numba的两个功能,以及它们如何与Dask组合:

  1. Numba的模具装饰工

  2. NumPy的广义通用函数

这最初是在此处发布为博客文章

Numba模具介绍

许多阵列计算功能仅在阵列的本地区域上运行。这在图像处理,信号处理,仿真,微分方程解,异常检测,时间序列分析等中很常见。通常,我们编写如下代码:

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) 
每个循环294毫秒±3.71毫秒(平均±标准偏差,共7次运行,每个循环1次) 

但是,如果我们使用Numba JIT编译此函数,则它运行得更快。

@numba.njit def smooth(x): return _smooth(x) %timeit smooth(x) 
每个回路70.3 µs±30.4 µs(平均±标准偏差,共运行7次,每个回路1个) 

对于那些计数,速度要快1000倍!

注意:该函数已经以``scipy.ndimage.uniform_filter''的形式存在,并且以相同的速度运行。

达斯克阵列

在这些应用程序中,人们经常有很多这样的数组,他们想将此功能应用于所有数组。原则上,他们可以使用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或current.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 
  数组
字节数 10.00 GB 50.00兆字节
形状 (10000、1000、1000) (50,1000,1000)
计数 200个任务 200块
类型 int8 numpy.ndarray
1000100010000

广义通用函数

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

NumPy文件: https //docs.scipy.org/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的2d数组并产生相同尺寸的int8的2d数组。

这种注释是一个很小的变化,但是它为Dask等其他系统提供了足够的信息来智能地使用它。除了调用像这样的函数外map_blocks,我们还可以直接使用该函数,就像我们的Dask数组只是一个非常大的NumPy数组一样。

# Before gufuncs y = x.map_blocks(smooth, dtype='int8') # After gufuncs y = smooth(x) y 
  数组
字节数 10.00 GB 50.00兆字节
形状 (10000、1000、1000) (50,1000,1000)
计数 800项任务 200块
类型 int8 numpy.ndarray
1000100010000

很好 如果您使用gufunc语义编写库代码,则该代码仅适用于Dask之类的系统,而无需为并行计算提供显式支持。这使用户的生活更加轻松。

启动Dask Client for Dashboard

启动Dask Client是可选的。它将启动仪表板,这有助于获得对计算的了解。

from dask.distributed import Client, progress client = Client(threads_per_worker=4, n_workers=1, processes=False, memory_limit='4GB') client 

客户

  • 工人数: 1
  • 核心数: 4
  • 内存: 4.00 GB
y.max().compute() 
122 

GPU版本

Numba还支持在兼容GPU设备上使用CUDA进行JIT编译。

使用numba.cuda.jit可以使单个V100 GPU上的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) 

完整的笔记本在这里

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章