改进移动窗口计算的内存消耗和速度

Question

在这种移动window计算中是否有可能获得更好的性能（内存消耗和速度）？我有一个 1000x1000 numpy 数组，我将 16x16 windows 遍历整个数组，最后对每个 window 应用一些函数（在这种情况下，离散余弦变换。）

import numpy as np
from scipy.fftpack import dct
from skimage.util import view_as_windows

X = np.arange(1000*1000, dtype=np.float32).reshape(1000,1000)
window_size = 16
windows = view_as_windows(X, (window_size,window_size))
dcts = np.zeros(windows.reshape(-1,window_size, window_size).shape, dtype=np.float32)
for idx, window in enumerate(windows.reshape(-1,window_size, window_size)):
    dcts[idx, :, :] = dct(window)
dcts = dcts.reshape(windows.shape)

此代码占用太多内存（在上面的示例中，内存消耗还算不错 - windows 使用 1Gb，dcts 也需要 1Gb）并且需要 25 秒才能完成。我有点不确定我做错了什么，因为这应该是一个简单的计算（例如过滤图像）。有没有更好的方法来完成这个？

更新：

我最初担心金顿解决方案生成的数组与我最初的方法有很大不同，但差异仅限于边界，因此不太可能对大多数应用程序造成严重问题。唯一剩下的问题是这两种解决方案都非常慢。目前，第一个解决方案需要 1 分钟 10 秒，第二个解决方案需要 59 秒。

更新 2：

我注意到到目前为止最大的罪魁祸首是 dct 和 np.mean。即使 generic_filter 使用 mean 的 "cythonized" 版本也有不错的表现（8.6 秒），但有瓶颈：

import bottleneck as bp
def func(window, shape):
    window = window.reshape(shape)
    #return np.abs(dct(dct(window, axis=1), axis=0)).mean()
    return bp.nanmean(dct(window))

result = scipy.ndimage.generic_filter(X, func, (16, 16),
                                      extra_arguments=([16, 16],))

我目前正在阅读如何使用 numpy 包装 C 代码以替换 scipy.fftpack.dct。如果有人知道该怎么做，我将不胜感激。

Answer 1

skimage.util.view_as_windows 正在使用跨步技巧来制作不使用任何额外内存的重叠 "windows" 数组。

但是，当您创建形状形状的新数组时，它需要的内存是原始 X 数组或 windows 数组所用内存的 ~32 倍 (16 x 16)。

根据您的评论，您的最终结果是 dcts.reshape(windows.shape).mean(axis=2).mean(axis=2) - 取每个 window 的 dct 的平均值。

因此，在循环内取均值而不存储 windows:

的巨大中间数组会更节省内存（尽管在性能方面相似）

import numpy as np
from scipy.fftpack import dct
from skimage.util import view_as_windows

X = np.arange(1000*1000, dtype=np.float32).reshape(1000,1000)
window_size = 16
windows = view_as_windows(X, (window_size, window_size))
dcts = np.zeros(windows.shape[:2], dtype=np.float32).ravel()
for idx, window in enumerate(windows.reshape(-1, window_size, window_size)):
    dcts[idx] = dct(window).mean()
dcts = dcts.reshape(windows.shape[:2])

另一种选择是scipy.ndimage.generic_filter。它不会增加太多性能（瓶颈是内循环中的 python 函数调用），但是您将有更多的边界条件选项，并且内存效率相当高：

import numpy as np
from scipy.fftpack import dct
import scipy.ndimage

X = np.arange(1000*1000, dtype=np.float32).reshape(1000,1000)

def func(window, shape):
    window = window.reshape(shape)
    return dct(window).mean()

result = scipy.ndimage.generic_filter(X, func, (16, 16),
                                      extra_arguments=([16, 16],))

Answer 2

由于scipy.fftpack.dct沿输入数组的最后一个轴计算单独的变换，您可以将循环替换为：

windows = view_as_windows(X, (window_size,window_size))
dcts = dct(windows)
result1 = dcts.mean(axis=(2,3))

现在只有 dcts 数组需要大量内存，而 windows 仍然只是 X 的一个视图。而且因为 DCT 是通过单个函数调用计算的，所以速度也快得多。但是，由于 windows 重叠，所以重复计算很多。这可以通过只为每个子行计算一次 DCT 来克服，然后是 windowed 平均值：

ws = window_size
row_dcts = dct(view_as_windows(X, (1, ws)))
cs = row_dcts.squeeze().sum(axis=-1).cumsum(axis=0)
result2 = np.vstack((cs[ws-1], cs[ws:]-cs[:-ws])) / ws**2

虽然在代码清晰度方面似乎失去了效率......但基本上这里的方法是首先计算 DCT，然后通过对 2D [=41] 求和来取 window 平均值=] 然后除以 window 中的元素数。 DCT 已经通过行向移动计算 windows，因此我们对这些 windows 进行常规求和。然而，我们需要对列进行移动 window 求和，以得出正确的 2D window 求和。为了有效地做到这一点，我们使用了 cumsum 技巧，其中：

sum(A[p:q])  # q-p == window_size

相当于：

cs = cumsum(A)
cs[q-1] - cs[p-1]

这避免了一遍又一遍地对完全相同的数字求和。不幸的是，它不适用于第一个 window（当 p == 0 时），因此我们必须只取 cs[q-1] 并将其与其他 window 和叠加在一起。最后我们除以元素的数量得到 2D window 平均值。

如果你喜欢做 2D DCT，那么第二种方法就没那么有趣了，因为你最终需要完整的 985 x 985 x 16 x 16 数组才能取平均值。

以上两种方法应该是等价的，但使用 64 位浮点数执行算术可能是个好主意：

np.allclose(result1, result2, atol=1e-6)
# False
np.allclose(result1, result2, atol=1e-5)
# True

改进移动窗口计算的内存消耗和速度

Improving moving-window computation in memory consumption and speed

python

performance

numpy