在多维数组中将滑动窗口设置为列(来自MATLAB的IM2COL)-Python

时间:2019-07-16 03:41:04

标签: python performance numpy vectorization

目前,我有一个4d数组

arr = np.arange(48).reshape((2,2,3,4))

我想将一个以2d数组作为输入的函数应用于从arr分割的每个2d数组。我已经搜索并阅读了this question,这正是我想要的。

我正在使用的功能是im2col_sliding_broadcasting(),我从here获得。它以一个2d数组和2个元素的列表作为输入,并返回一个2d数组。就我而言:它需要3x4 2d数组和列表[2, 2]并返回4x6 2d数组。

我考虑使用apply_along_axis(),但是as said仅接受1d函数作为参数。我无法以这种方式应用im2col函数。

我想要一个形状为2x2x4x6的输出。我当然可以通过for循环来实现这一点,但是我听说这太浪费时间了:

import numpy as np

def im2col_sliding_broadcasting(A, BSZ, stepsize=1):
    # source: https://stackoverflow.com/a/30110497/10666066
    # Parameters
    M, N = A.shape
    col_extent = N - BSZ[1] + 1
    row_extent = M - BSZ[0] + 1

    # Get Starting block indices
    start_idx = np.arange(BSZ[0])[:, None]*N + np.arange(BSZ[1])

    # Get offsetted indices across the height and width of input array
    offset_idx = np.arange(row_extent)[:, None]*N + np.arange(col_extent)

    # Get all actual indices & index into input array for final output
    return np.take(A, start_idx.ravel()[:, None] + offset_idx.ravel()[::stepsize])

arr = np.arange(48).reshape((2,2,3,4))
output = np.empty([2,2,4,6])

for i in range(2):
    for j in range(2):
        temp = im2col_sliding_broadcasting(arr[i, j], [2,2])
        output[i, j] = temp

因为我的arr实际上是一个10000x3x64x64数组。所以我的问题是:还有另一种方法可以更有效地做到这一点吗?

1 个答案:

答案 0 :(得分:1)

我们可以利用基于np.lib.stride_tricks.as_stridedscikit-image's view_as_windows来获取滑动窗口。 More info on use of as_strided based view_as_windows

from skimage.util.shape import view_as_windows

W1,W2 = 2,2 # window size

# create sliding windows along last two axes1
w = view_as_windows(arr,(1,1,W1,W2))[...,0,0,:,:]

# Merge the window axes (tha last two axes) and 
# merge the axes along which those windows were created (3rd and 4th axes)
outshp = arr.shape[:-2] + (W1*W2,) + ((arr.shape[-2]-W1+1)*(arr.shape[-1]-W2+1),)
out = w.transpose(0,1,4,5,2,3).reshape(outshp)

最后一步强制复制。因此,请尽可能跳过它。