将3D numpy数组高效转换为1D numpy数组

时间:2015-06-16 22:41:03

标签: python arrays numpy

我有一个这种形式的3D numpy数组:

>>>img.shape
(4504932, 2, 2)

>>> img
array([[[15114, 15306],
    [15305, 15304]],

   [[15305, 15306],
    [15303, 15304]],

   [[15305, 15306],
    [15303, 15304]],

   ..., 

   [[15305, 15302],
    [15305, 15302]]], dtype=uint16)

我要将其转换为1D numpy数组,其中每个条目是上述img numpy数组中每个2x2子矩阵的总和。

我已经能够使用以下方式完成此任务:

img_new = np.array([i.sum() for i in img])
>>> img_new
array([61029, 61218, 61218, ..., 61214, 61214, 61214], dtype=uint64)

这正是我想要的。但这太慢了(大约需要10秒钟)。有没有更快的方法可以使用?我在上面img.shape中包含了这个numpy数组的大小。

编辑 - 附加信息: 我的img数组也可以是4x4,5x5,7x7等子矩阵形式的3D数组。这由变量sub_rowssub_cols指定。

3 个答案:

答案 0 :(得分:4)

img.sum(axis=(1, 2))

sum允许您指定要求和的轴或轴,而不仅仅是对整个数组求和。这允许NumPy在C中循环遍历数组,并且每次执行只执行一些机器指令,而不是必须通过Python字节码评估循环并创建大量的包装器对象以粘贴在列表中。

答案 1 :(得分:0)

使用numpy方法(apply_over_axes)通常更快,事实就是如此。我刚刚测试了4000x2x2阵列:

img = np.random.rand(4000,2,2)
timeit(np.apply_along_axis(np.sum, img, [1,2]))
# 1000 loops, best of 3: 721 us per loop
timeit(np.array([i.sum() for i in img]))
# 100 loops, best of 3: 17.2 ms per loop

答案 2 :(得分:0)

您可以使用np.einsum -

img_new = np.einsum('ijk->i',img)

验证结果

In [42]: np.array_equal(np.array([i.sum() for i in img]),np.einsum('ijk->i',img))
Out[42]: True

运行时测试

In [34]: img = np.random.randint(0,10000,(10000,2,2)).astype('uint16')

In [35]: %timeit np.array([i.sum() for i in img]) # Original approach
10 loops, best of 3: 92.4 ms per loop

In [36]: %timeit img.sum(axis=(1, 2)) # From other solution
1000 loops, best of 3: 297 µs per loop

In [37]: %timeit np.einsum('ijk->i',img)
10000 loops, best of 3: 102 µs per loop
相关问题