Python实现im2col和col2im函数完整代码

在这里贴上了关于im2col和col2im函数的代码，已覆盖了基本功能，供大家调用学习。

实现im2col函数：

import numpy as np


def im2col(img, kernel_size, strides=1, paddings=0, dilations=1):
    if not isinstance(kernel_size, list):
        k = kernel_size
        kernel_size = [k, k]
    k_h, k_w = kernel_size

    if isinstance(strides, list):
        s_h, s_w = strides
    else:
        s_h = s_w = strides

    if isinstance(paddings, list):
        if len(paddings) == 4:
            ps = paddings
        else:
            p_h, p_w = paddings
            ps = [p_h, p_w, p_h, p_w]
    else:
        p = paddings
        ps = [p, p, p, p]

    if not isinstance(dilations, list):
        d = dilations
        dilations = [d, d]
    d_h, d_w = dilations

    img = np.pad(img, ((0, 0), (0, 0), (ps[0], ps[2]), (ps[1], ps[3])), 'constant')

    N, C, H, W = img.shape

    dk_h = (d_h * (k_h - 1) + 1)
    dk_w = (d_w * (k_w - 1) + 1)
    out_h = (H - dk_h) // s_h + 1
    out_w = (W - dk_w) // s_w + 1

    outsize = out_w * out_h
    col = np.empty((N, C, k_w * k_h, outsize))

    for y in range(out_h):
        y_start = y * s_h
        y_end = y_start + dk_h
        for x in range(out_w):
            x_start = x * s_w
            x_end = x_start + dk_w
            col[:, :, 0:, y * out_w + x] = img[:, :, y_start:y_end:d_h, x_start:x_end:d_w].reshape([N, C, k_h * k_w])
    return paddle.to_tensor(col.reshape([N, -1, outsize]))


import paddle

x = paddle.to_tensor([[48, 38, 38, 59, 38],
                      [38, 11, 25, 52, 44],
                      [60, 69, 49, 93, 66],
                      [88, 8, 47, 14, 47],
                      [96, 37, 56, 86, 54]], dtype='float32'
                     ).unsqueeze(0).unsqueeze(0)

result = im2col(x, [3, 3], strides=1, paddings=0, dilations=2)
print(result)

实现col2im函数：

import numpy as np


def col2im(mtx, image_size, kernel_size, strides=1, paddings=0, dilations=1):
    if not isinstance(kernel_size, list):
        k = kernel_size
        kernel_size = [k, k]
    k_h, k_w = kernel_size

    if isinstance(strides, list):
        s_h, s_w = strides
    else:
        s_h = s_w = strides

    if isinstance(paddings, list):
        if len(paddings) == 4:
            ps = paddings
        else:
            p_h, p_w = paddings
            ps = [p_h, p_w, p_h, p_w]
    else:
        p = paddings
        ps = [p, p, p, p]

    if isinstance(image_size, list):
        image_h, image_w = image_size
    else:
        image_h = image_w = image_size

    if not isinstance(dilations, list):
        d = dilations
        dilations = [d, d]
    d_h, d_w = dilations

    N, kC, _ = mtx.shape
    r_C = kC // (k_h * k_w)
    img = np.zeros([N, r_C, image_h, image_w])

    result = np.pad(img, ((0, 0), (0, 0), (ps[0], ps[2]), (ps[1], ps[3])), 'constant')
    weight = np.pad(img, ((0, 0), (0, 0), (ps[0], ps[2]), (ps[1], ps[3])), 'constant')

    dk_h = (d_h * (k_h - 1) + 1)
    dk_w = (d_w * (k_w - 1) + 1)

    out_h = (result.shape[-2] - dk_h) // s_h + 1
    out_w = (result.shape[-1] - dk_w) // s_w + 1

    col = 0
    for i in range(out_h):
        for j in range(out_w):
            result[:, :, i * s_h:i * s_h + dk_h:d_h, j * s_w:j * s_w + dk_w:d_w] += \
                mtx[:, :, col].reshape([N, r_C, k_h, k_w]).numpy()
            weight[:, :, i * s_h:i * s_h + dk_h:d_h, j * s_w:j * s_w + dk_w:d_w] += \
                np.ones([N, r_C, k_h, k_w])
            col += 1

    with np.errstate(divide='ignore', invalid='ignore'):
        result_avg = result / weight
        result_avg[~ np.isfinite(result_avg)] = 0

    return paddle.to_tensor(result_avg[:, :, ps[0]:image_h + ps[0], ps[1]:image_w + ps[1]])


import paddle

x = paddle.to_tensor(
    [[[48.],
      [38.],
      [38.],
      [60.],
      [49.],
      [66.],
      [96.],
      [56.],
      [54.]]])

result = col2im(x, [5, 5], [3, 3], strides=1, paddings=0, dilations=2)
print(result)