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

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 optimizations 作为 vectorize a lot of operations 的处理器)，所以它比迭代快得多。代码中另一个耗时的操作是附加到 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 数组!使用更大的数组，您可以期望看到更大的加速。