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

import numpy as np

def im2col(input_data, filter_h, filter_w, stride=1, pad=0):
    """
    将多维数组转换为二维矩阵，并按列排列滑动窗口中的元素
    :param input_data: 输入数据，形状为(N, C, H, W)
    :param filter_h: 滑动窗口的高度
    :param filter_w: 滑动窗口的宽度
    :param stride: 步幅
    :param pad: 填充
    :return: 二维矩阵，形状为(N*OH*OW, C*FH*FW)，其中OH和OW为输出数据的高度和宽度
    """
    N, C, H, W = input_data.shape
    out_h = (H + 2*pad - filter_h) // stride + 1
    out_w = (W + 2*pad - filter_w) // stride + 1

    img = np.pad(input_data, [(0,0), (0,0), (pad,pad), (pad,pad)], 'constant')
    col = np.zeros((N, C, filter_h, filter_w, out_h, out_w))

    for y in range(filter_h):
        y_max = y + stride*out_h
        for x in range(filter_w):
            x_max = x + stride*out_w
            col[:, :, y, x, :, :] = img[:, :, y:y_max:stride, x:x_max:stride]

    col = col.transpose(0, 4, 5, 1, 2, 3).reshape(N*out_h*out_w, -1)
    return col