使用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)
完整的筆記本在這裏