python - 整个块的 dask.array 缓存

Question

我有一些大型 4 维数据集 (100 GB+)，我使用h5py+dask.array来处理和可视化它们。例如形状(1024, 1024, 256, 256)，块大小为(32, 32, 32, 32)。由于数据通常保存在非 SSD 硬盘上并经过压缩，因此需要一些时间才能将 HDF5 文件中的整个块放入 NumPy 数组中。

在探索数据时，能够快速浏览数据集非常好。这方面的一个例子，在一个小的人工 HDF5 数据集上：

import dask.array as da
dask_array = da.zeros((100, 100, 100, 100), chunks=(50, 50, 100, 100), dtype="uint32")
dask_array.to_hdf5("test_data.hdf5", "/test_data", compression="lzf")

现在，假设我们想通过抓取前两个维度中的单个位置以及后两个维度中的完整图像来探索数据集。通过使用dask的机会缓存

from time import time
import numpy as np
import h5py
import dask.array as da
from dask.cache import Cache
cache = Cache(1e9)
cache.register()

dask_array = da.from_array(h5py.File("test_data.hdf5")["/test_data"], chunks=(50, 50, 100, 100))

t0 = time()
np.array(dask_array[5, 2]) # Takes 0.25 seconds
print(time() - t0)

t0 = time()
# Grab a different index, which in the same chunk/block as the previous
np.array(dask_array[15, 2]) # Takes 0.25 seconds
print(time() - t0)

t0 = time()
# Grab the same index as the first one
np.array(dask_array[5, 2]) # Takes 0.002 seconds
print(time() - t0)

因此，机会缓存存储特定位置，而不是整个块。在数据集上导航时，这会引入延迟，因为必须从硬盘驱动器中为块中的每个新位置加载整个块。有没有办法dask通过机会缓存自动缓存整个块？

---

我猜它可以使用dask_array.persist，像这样：

dask_chunk = dask_array[0:50, 0:50].persist()
np.array(dask_chunk[5, 2])

但是，这需要一些额外的“基础设施”来跟踪dask_chunk. 因此，一些块/块范围的机会缓存会更好。

标签： pythondask