python – 按最大大小将numpy数组拆分为块

我有一些非常大的二维numpy数组.一个数据集是55732乘257659,超过140亿个元素.因为我需要执行的一些操作抛出MemoryErrors,所以我想尝试将数组拆分成一定大小的块并对着块运行它们. (我可以在每个部分运行操作后聚合结果.)我的问题是MemoryErrors这一事实意味着我可以以某种方式限制数组的大小,而不是将它们分成不变的数量.

举个例子,让我们生成一个1009乘1009的随机数组:

a = numpy.random.choice([1,2,3,4], (1009,1009))

我的数据不一定可以均匀分割,并且绝对不能保证按我想要的大小分割.所以我选择1009因为它是素数.

让我们也说我想要它们的块数不超过50比50.因为这只是为了避免极大数组的错误,如果结果不准确就没关系.

如何将其拆分为所需的块?

我正在使用Python 3.6 64位和numpy 1.14.3(最新版).

有关

我见过this function that uses reshape,但是如果行数和列数没有完全划分大小,它就不起作用.

This question(以及其他类似的)有答案解释如何分割成一定数量的块,但这并不能解释如何分割成一定的大小.

我也看过this question,因为这实际上是我的确切问题.答案和评论建议切换到64位(我已经拥有)并使用numpy.memmap.没有帮助.

最佳答案
这可以这样做,使得所得到的阵列具有略小于所需最大值的形状,或者使得它们具有恰好所需的最大值,除了末端的一些剩余部分.

基本逻辑是计算用于拆分阵列的参数,然后使用array_split沿阵列的每个轴(或维度)拆分阵列.

我们需要numpy和math模块以及示例数组:

import math
import numpy

a = numpy.random.choice([1,2,3,4], (1009,1009))

略低于最大值

逻辑

首先沿着要在元组中将其拆分的每个维度存储最终块大小的形状:

chunk_shape = (50, 50)

array_split仅一次沿一个轴(或维度)或一个数组分割.所以让我们从第一个轴开始吧.

>计算将数组拆分为以下所需的部分数:

num_sections = math.ceil(a.shape[0] / chunk_shape[0])

在我们的示例中,这是21(1009/50 = 20.18).
>现在拆分它:

first_split = numpy.array_split(a, num_sections, axis=0)

这给了我们一个21(已请求部分的数量)numpy数组的列表,这些数组被分割,因此它们在第一维中不大于50:

print(len(first_split))
# 21
print({i.shape for i in first_split})
# {(48, 1009), (49, 1009)}
# These are the distinct shapes, so we don't see all 21 separately

在这种情况下,它们沿着该轴是48和49.
>对于第二个维度,我们可以对每个新数组执行相同的操作:

num_sections = math.ceil(a.shape[1] / chunk_shape[1])
second_split = [numpy.array_split(a2, num_sections, axis=1) for a2 in first_split]

这给了我们一个列表清单.每个子列表包含我们想要的大小的numpy数组:

print(len(second_split))
# 21
print({len(i) for i in second_split})
# {21}
# All sublists are 21 long
print({i2.shape for i in second_split for i2 in i})
# {(48, 49), (49, 48), (48, 48), (49, 49)}
# Distinct shapes

功能齐全

我们可以使用递归函数为任意维度实现这个:

def split_to_approx_shape(a, chunk_shape, start_axis=0):
    if len(chunk_shape) != len(a.shape):
        raise ValueError('chunk length does not match array number of axes')

    if start_axis == len(a.shape):
        return a

    num_sections = math.ceil(a.shape[start_axis] / chunk_shape[start_axis])
    split = numpy.array_split(a, num_sections, axis=start_axis)
    return [split_to_approx_shape(split_a, chunk_shape, start_axis + 1) for split_a in split]

我们称之为:

full_split = split_to_approx_shape(a, (50,50))
print({i2.shape for i in full_split for i2 in i})
# {(48, 49), (49, 48), (48, 48), (49, 49)}
# Distinct shapes

精确的形状加上余数

逻辑

如果我们想成为一个有点发烧友,并且所有新数组都是指定大小,除了尾随剩余数组,我们可以通过传递一个索引列表来分割到array_split.

>首先构建索引数组:

axis = 0
split_indices = [chunk_shape[axis]*(i+1) for i  in range(math.floor(a.shape[axis] / chunk_shape[axis]))]

这使用了一个索引列表,每个索引从最后一个50:

print(split_indices)
# [50, 100, 150, 200, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700, 750, 800, 850, 900, 950, 1000]

>然后拆分:

first_split = numpy.array_split(a, split_indices, axis=0)
print(len(first_split))
# 21
print({i.shape for i in first_split})
# {(9, 1009), (50, 1009)}
# Distinct shapes, so we don't see all 21 separately
print((first_split[0].shape, first_split[1].shape, '...', first_split[-2].shape, first_split[-1].shape))
# ((50, 1009), (50, 1009), '...', (50, 1009), (9, 1009))

>然后再次为第二轴:

axis = 1
split_indices = [chunk_shape[axis]*(i+1) for i  in range(math.floor(a.shape[axis] / chunk_shape[axis]))]
second_split = [numpy.array_split(a2, split_indices, axis=1) for a2 in first_split]
print({i2.shape for i in second_split for i2 in i})
# {(9, 50), (9, 9), (50, 9), (50, 50)}

功能齐全

调整递归函数:

def split_to_shape(a, chunk_shape, start_axis=0):
    if len(chunk_shape) != len(a.shape):
        raise ValueError('chunk length does not match array number of axes')

    if start_axis == len(a.shape):
        return a

    split_indices = [
        chunk_shape[start_axis]*(i+1)
        for i in range(math.floor(a.shape[start_axis] / chunk_shape[start_axis]))
    ]
    split = numpy.array_split(a, split_indices, axis=start_axis)
    return [split_to_shape(split_a, chunk_shape, start_axis + 1) for split_a in split]

我们称之为完全相同的方式:

full_split = split_to_shape(a, (50,50))
print({i2.shape for i in full_split for i2 in i})
# {(9, 50), (9, 9), (50, 9), (50, 50)}
# Distinct shapes

额外说明

性能

这些功能看起来非常快.我能够将我的示例数组(超过140亿个元素)分成1000个1000个形状的片段(产生超过14000个新阵列),在0.05秒内完成以下任一功能:

print('Building test array')
a = numpy.random.randint(4, size=(55000, 250000), dtype='uint8')
chunks = (1000, 1000)
numtests = 1000
print('Running {} tests'.format(numtests))
print('split_to_approx_shape: {} seconds'.format(timeit.timeit(lambda: split_to_approx_shape(a, chunks), number=numtests) / numtests))
print('split_to_shape: {} seconds'.format(timeit.timeit(lambda: split_to_shape(a, chunks), number=numtests) / numtests))

输出:

Building test array
Running 1000 tests
split_to_approx_shape: 0.035109398348040485 seconds
split_to_shape: 0.03113800323300747 seconds

我没有用更高维度的数组测试速度.

形状小于最大值

如果任何尺寸的尺寸小于指定的最大值,这些功能都可以正常工作.这不需要特殊的逻辑.

转载注明原文:python – 按最大大小将numpy数组拆分为块 - 代码日志