From 6be2b1bfa265932b0d6dda5035097ab5dc1ccd53 Mon Sep 17 00:00:00 2001 From: Blake Griffith Date: Wed, 6 May 2015 15:35:36 -0500 Subject: [PATCH] Smarter chunk selection. --- skimage/util/process.py | 27 ++++++++++++++++++++++----- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/skimage/util/process.py b/skimage/util/process.py index c03b1c5b..177651ff 100644 --- a/skimage/util/process.py +++ b/skimage/util/process.py @@ -1,3 +1,4 @@ +from math import ceil from multiprocessing import cpu_count import dask.array as da @@ -11,19 +12,35 @@ def _get_chunks(shape, ncpu): available processors. The last chunk in each dimension absorbs the remainder array elements if the number of cpus does not divide evenly into the number of array elements. + + >>> _get_chunks((4, 4), 4) + ((2, 2), (2, 2)) + >>> _get_chunks((4, 4), 2) + ((2, 2), (4,)) + >>> _get_chunks((5, 5), 2) + ((2, 3), (5,)) + >>> _get_chunks((2, 4), 2) + ((1, 1), (4,)) """ chunks = [] + nchunks_per_dim = int(ceil(ncpu ** (1./len(shape)))) + used_chunks = 1 for i in shape: - regular_chunk = i // ncpu - remainder_chunk = regular_chunk + (i % ncpu) + if used_chunks < ncpu: + regular_chunk = i // nchunks_per_dim + remainder_chunk = regular_chunk + (i % nchunks_per_dim) - if regular_chunk == 0: - chunk_lens = (remainder_chunk,) + if regular_chunk == 0: + chunk_lens = (remainder_chunk,) + else: + chunk_lens = ((regular_chunk,) * (nchunks_per_dim - 1) + + (remainder_chunk,)) else: - chunk_lens = (regular_chunk,) * (ncpu - 1) + (remainder_chunk,) + chunk_lens = (i,) chunks.append(chunk_lens) + used_chunks *= nchunks_per_dim return tuple(chunks)