基于 Python tile 的 numpy 数组处理

Question

我有一个非常大的 numpy 数组，大小为 [256,256,256,256]，它占用了大约 8GB 的​​内存。

我想使用 multiprocessing 快速处理数据，类似于基于 tile 的渲染软件（例如 blender）使用的方法。

我想将我的数据分成更小的块，为它们创建一个进程列表，当一个完成后，开始处理下一个。我当前的方法使用 V 拆分，然后循环执行 hsplit 并获取块。下面是一个小得多的数据集的代码示例（9 x 9 分为 3 x 3 块）：

import numpy as np

data = np.array(np.arange(0,81)).reshape((9,9))

print(data.shape)
print(data)
print("#"*30)
data = np.array([np.vsplit(set,3) for set in np.hsplit(data,3)])
print(data.shape)
print(data)
print("#"*30)

这并不完全符合我的要求，它创建了 9 个 3 x 3 的块，但这是一个小问题。主要问题是我不认为通过这样的数组循环来应用 vsplit 非常有效，但我不知道有一个内置函数可以自动执行此操作而无需循环。

我尝试使用 map 来查看它是否以更有效的方式自动应用 vsplit，但计时结果非常相似，所以我认为不是这样。有关如何优化此问题的任何建议？

Answer 1

您可以使用SciKit-Image's skimage.util.view_as_blocks()，它可以在不移动或复制任何数据的情况下解决您的阻塞问题。

skimage.util.view_as_blocks(data, (3, 3))

Answer 2

您想要基于as_strided 的东西。我推荐my recipe here。使用这个你可以做到：

patches = window_nd(data, window = 3, steps = 3)

要传递到多处理，您需要一个生成器。比如：

def patch_gen(data, window, **kwargs):
    dims = len(data.shape)
    patches = window_nd(data, window, **kwargs)
    patches = patches.reshape((-1,) + patches.shape[-dims:])
    for i in range(patches.shape[0]):
        yield patches[i]

您也可以使用来自@NilsWerner 的view_as_blocks：

def patch_gen(data, window_shape):
    dims = len(data.shape)
    patches = skimage.util.view_as_blocks(data, window_shape)
    patches = patches.reshape((-1,) + patches.shape[-dims:])
    for i in range(patches.shape[0]):
        yield patches[i]

然后你可以这样做：

with Pool(processors) as p:
    out = p.map(func, patch_gen(data, window = 3, steps = 3))

这种方法应该只使用数组的原始内存，而不是在任何时候复制（我认为）。这样您就不会浪费时间和内存来复制原始数据，所有结果都只是指向原始数据的指针。

警告：不要写入补丁！至少，in general it's a bad idea to write to as_strided results. 如果你没有重叠的窗口，你可能没问题，但是一旦你的窗口重叠（即stepswindow）你就会一团糟你的手。