python - 使用 2D 掩码和整个矩阵运算索引的 3D 或 4D Numpy 数组

Question

我试图对整个 3D 或 4D 数组进行一些操作，但仅限于较小尺寸的子组(包含在较大数组中的 2D 数组)。

示例:

input = np.arange(75).reshape((3, 5, 5))  # or any other 3D or 4D matrix. 
mask_hor = np.arange(-1, 2)
mask_ver = mask_hor[:, None]

output = np.zeros((3, 3, 3))

for i in range(1, 5):
  for j in range(1, 5):
    output[:, i-1, j-1] = foo(input[:, i+mask_ver, j+mask_hor])

其中 foo 是对输入的某种操作

我的问题是: 是否有一个方法/掩码可以传递给输入，以便我可以摆脱嵌套的 for 循环？我主要是在寻找加速。

谢谢!

Answer 1

这更像是一种快速而肮脏的优化，而不是任何优雅的优化。为了便于论证，我们将对窗口中的 9 个元素求和作为我们的 foo 函数。

import numpy as np
from scipy import ndimage

# take the sum of a 3x3 window of a matrix
def foo_lin_mat(mat):
    return mat.sum(axis=(-2, -1))  # sum over the last two axes

# sum up the individual matrices
def foo_lin_nine(m1, m2, m3, m4, m5, m6, m7, m8, m9):
    return m1 + m2 + m3 + m4 + m5 + m6 + m7 + m8 + m9

# compute foo on an input matrix by shifting the mask around
def nestfor(input, foo):
    depth, n, m = input.shape
    output = np.zeros((depth, n - 2, m - 2))
    mask_hor = np.arange(-1, 2)
    mask_ver = mask_hor[:, None]
    for i in range(1, n - 1):
        for j in range(1, m - 1):
            output[:, i - 1, j - 1] = foo(input[:, i + mask_ver, j + mask_hor])
    return output

# compute foo on an input matrix by breaking the input matrix into 9 submatrices 
def flatargs(input, foo):
    depth, n, m = input.shape
    return foo(input[:, :n-2, :m-2], 
               input[:, 1:n-1, :m-2],
               input[:, 2:, :m-2],
               input[:, :n-2, 1:m-1],
               input[:, 1:n-1, 1:m-1],
               input[:, 2:, 1:m-1],
               input[:, :n-2, 2:],
               input[:, 1:n-1, 2:],
               input[:, 2:, 2:], )

# compute the sum of a window using ndimage.convolve
def convolve(input, mask):
    mask = np.ones((1, 3, 3))
    out = ndimage.convolve(input, mask)
    # cut off the outer edges
    return out[1:-1, 1:-1]

所以我们有三个函数，它们将采用一个矩阵并对各个 3x3 窗口求和。我已经确认他们最后吐出了相同的矩阵。至于基准测试

In [62]: %timeit  nestfor(input, foo_lin_mat)
1000 loops, best of 3: 261 µs per loop

In [63]: %timeit flatargs(input, foo_lin_nine)
10000 loops, best of 3: 35.8 µs per loop

In [66]: mask = np.ones((1,3,3))

In [69]: %timeit convolve(input, mask)
The slowest run took 6.12 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 3: 42.2 µs per loop

即flatargs 版本比原始嵌套 for 循环版本快约 7 倍。

如果您的 foo 函数是输入窗口的线性函数，您还可以使用 ndimage.convolve 函数来进行窗口化，如 convolve 中所示 函数在这里。阅读最终代码可能会更容易一些，但您必须小心用于掩码的数组。