使用Numba进行模具计算
这款笔记本结合Numba,高性能的Python编译器,用DASK阵列。
特别是,我们展示了Numba的两个功能,以及它们如何与Dask组合:
这最初是在此处发布为博客文章
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
|
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
|
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
客户
|
簇
|
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)
完整的笔记本在这里