python - 如何向量化多维 numpy 数组中的提取

Question

我正在努力解决矢量化问题，乍一看似乎很简单:

假设我有 100 张大小为 (7, 7) 的图像，其中 2 个 channel 由 numpy 表示大小为 (100, 2, 7, 7) 的数组。我想在所有图像上提取小补丁(比如说大小(2,3,3))，但这些补丁并不位于每个图像上的同一位置。补丁的位置由大小为 (2, 100) 的矩阵描述(每个图像一个 x 和一个 y)。 我可以通过对所有图像进行 for 循环来做到这一点，但这需要时间。

这里是一个示例代码:

data = np.arange(9800).reshape((100, 2, 7, 7))
size = 3
pos = np.random.randint(0, 7-size, (2, 100))
for i in range(a.shape[0]):
    patch = data[i, :, pos[0,i]:(pos[0,i]+size), pos[1,i]:pos[1,i]+size] 

换句话说，我想在没有 for 循环的情况下重现这段代码。 有人知道这样做吗？

Answer 1

numpy.lib.stride_tricks.as_strided 的标准技巧应该在这里工作:

das = np.lib.stride_tricks.as_strided(data, (5, 5, 100, 2, 3, 3), data.strides[-2:] + data.strides)
patches = das[(*pos, np.arange(100))]

您可以使用以下方法验证其是否有效:

for i in range(data.shape[0]):
    assert np.all(patches[i]==data[i,:,pos[0,i]:pos[0,i]+size,pos[1,i]:pos[1,i]+size])