Python tile based processing of numpy array

Question

I have a very large numpy array of size [256,256,256,256] which takes up about 8GB of ram.

I want to process the data quickly using multiprocessing, similar to the method used by tile based rendering software such as blender.

I would like to split my data into chunks of smaller size, create a list of processes for them and when one finishes, start working on the next. My current method uses a V split, which is then looped over to perform hsplit and get the chunks. Here is an example of the code for a much smaller data set (9 by 9 divided into 3 by 3 chunks):

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)

this doesn't quite do what I want it to, which is creating 9 chunks of 3 by 3 but that's a minor issue. The major issue is I don't think looping through an array like that to apply vsplit is very efficient, but I'm not aware of a built in function that does this automatically without looping.

I tried using map to see if it automatically applies vsplit in a more efficient way, but the timing results were very similar, so I don't think that's the case. Any suggestions on how to optimize this problem?

Answer 1

You want something as_strided based. I'd recommend my recipe here . Using this you can do:

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

To pass to multiprocessing, you'll need a generator. Something like:

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]

You can also use view_as_blocks as from @NilsWerner:

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]

Then you can do something like:

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

This method should only use the original memory of the array and not make a copy at any point (I think). That way you don't waste time and memory making copies of the original data, all the results are just pointers to the original data.

Caution: don't write to the patches, At least, in general it's a bad idea to write to as_strided results. If you don't have overlapping windows you might be ok, but as soon as your windows overlap (ie steps < window ) you'll have a mess on your hands.

Answer 2

You can use SciKit-Image's skimage.util.view_as_blocks() , which solves your blocking problem without moving or copying any of the data.

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