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

Question

我正在尝试找到一种方法来跨多个 2D 数组对每个元素执行操作，而不必遍历它们。或者至少，不需要两个 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 optimizations 为您的处理器提供 vectorize a lot of operations），它比迭代快得多。代码中另一个耗时的操作是追加到stddev_arr。附加到 numpy 数组是 slow 因为在添加新元素之前将 整个数组 复制到新内存中。现在您已经知道该数组需要多大了，所以您不妨预先分配它。

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)