从 Dask 数组中选择多个切片的最佳方法

Question

我正在使用 Dask 阵列生成分布在多个节点之间的大型 (65k x 65k x 3) 3D 信号。 在下一步中，我需要使用存储在 Dask 包中的切片从该数组中提取数千个切片。我的代码如下所示：

import dask.array as da
import dask.bag as db
from dask.distributed import Client

def pick_tile(window, signal):
    return np.array(surface[window])

def computation_on_tile(signal_tile):
   # do some rather short computation on a (n x n x 3) signal tile.

dask_client = Client(....)

signal_array = generate_signal(...)  # returns a dask array
signal_slices = db.from_sequence(generate_slices(...))  # fixed size slices

signal_tiles = signal_slices.map(pick_tile, signal=signal_array)

result = dask_client.compute(signal_tile.map(computation_on_tile), sync=True)

我的问题是计算需要很多时间。我尝试使用以下方法分散我的信号数组：

signal_array = dask_client.scatter(generate_signal(...))

但这对性能没有帮助（大约需要 12 分钟来计算）。相比之下，全信号和第一层的stdev的计算大约需要2分钟。

有没有一种从分布式 Dask 数组中挑选大量切片的有效方法？

Answer 1

如果您只有几千个切片，那么我建议使用普通的 Python 列表而不是 Dask Bag。它可能会更快，更简单。

然后你可以多次切片你的数组：

tiles = [dask_array[slc] for slc in slices]

如果你愿意，可以计算这些

tiles = dask.compute(*tiles)