迭代多维数组的有效方法？

Question

我试图找到一种方法来对跨多个二维数组的每个元素执行操作，而不必遍历它们。 或者至少，不需要两个for循环。 我的代码计算一系列图像（数组）上每个像素的标准偏差。 现在，图像的数量不是问题，而是数组的大小，使代码变得非常慢。 以下是我所拥有的工作示例。

import numpy as np

# reshape(# of image (arrays),# of rows, # of cols) 
a = np.arange(32).reshape(2,4,4)

stddev_arr = np.array([])
for i in range(4):
    for j in range(4): 
        pixel = a[0:,i,j]
        stddev = np.std(pixel) 
        stddev_arr = np.append(stddev_arr, stddev)

我的实际数据是 2000x2000，使这段代码循环了 4000000 次。 有一个更好的方法吗？ 任何建议都非常感谢。

Answer 1

您已经在使用numpy 。 numpy 的std()函数接受一个axis参数，该参数告诉它您希望它在哪个轴上操作（在这种情况下是第零轴）。 由于这种卸载的计算numpy的的C-后端（和可能使用SIMD优化您的处理器是矢量化很多操作），它是如此比迭代要快很多。 代码中另一个耗时的操作是附加到stddev_arr 。 追加到 numpy 数组很慢，因为在添加新元素之前将整个数组复制到新内存中。 现在您已经知道该数组需要多大，因此您不妨预先分配它。

a = np.arange(32).reshape(2, 4, 4)
stdev = np.std(a, axis=0)

这给出了一个4x4数组

array([[8., 8., 8., 8.],
       [8., 8., 8., 8.],
       [8., 8., 8., 8.],
       [8., 8., 8., 8.]])

要将其展平为一维数组，请执行flat_stdev = stdev.flatten() 。

比较执行时间：

# Using only numpy
def fun1(arr):
    return np.std(arr, axis=0).flatten()

# Your function
def fun2(arr):
    stddev_arr = np.array([])
    for i in range(arr.shape[1]):
        for j in range(arr.shape[2]): 
            pixel = arr[0:,i,j]
            stddev = np.std(pixel) 
            stddev_arr = np.append(stddev_arr, stddev)
    return stddev_arr


# Your function, but pre-allocating stddev_arr
def fun3(arr):
    stddev_arr = np.zeros((arr.shape[1] * arr.shape[2],))
    x = 0
    for i in range(arr.shape[1]):
        for j in range(arr.shape[2]): 
            pixel = arr[0:,i,j]
            stddev = np.std(pixel) 
            stddev_arr[x] = stddev
            x += 1
    return stddev_arr

首先，让我们确保所有这些函数都是等价的：

a = np.random.random((3, 10, 10))
assert np.all(fun1(a) == fun2(a))
assert np.all(fun1(a) == fun3(a))

是的，都给出相同的结果。 现在，让我们尝试使用更大的数组。

a = np.random.random((3, 100, 100))

x = timeit.timeit('fun1(a)', setup='from __main__ import fun1, a', number=10)
# x: 0.003302899989648722

y = timeit.timeit('fun2(a)', setup='from __main__ import fun2, a', number=10)
# y: 5.495519500007504

z = timeit.timeit('fun3(a)', setup='from __main__ import fun3, a', number=10)
# z: 3.6250679999939166

哇！ 仅通过预分配，我们就获得了约 1.5 倍的加速。 更令人惊叹的是：将 numpy 的std()与axis参数一起使用可提供 > 1000 倍的加速，而这仅适用于100x100数组！ 使用更大的数组，您可以期望看到更大的加速。

Answer 2

因此，根据您提供的内容，您可以以另一种方式重塑数组以对其进行矢量化以替换您的两个循环。 然后你只需要在你想要的轴上使用np.std一次。

a = np.arange(32).reshape(2, 4, 4)

a = a.reshape(2, -1).transpose()

stddev_arr = np.std(a, axis=1)