卷积操作

题目描述

实现一个多通道输入、单输出通道的二维卷积操作。给定一个形状为 的输入张量和形状为 的卷积核,以及步长 stride 和填充 padding,计算卷积输出。

思路

这是一道模拟题,按照卷积的定义逐步实现即可:

  1. 零填充:在输入的每个通道四周填充 padding 层的零,填充后尺寸变为
  1. 滑动窗口:以 stride 为步长,将 大小的窗口在填充后的输入上滑动。输出尺寸为:

$$

$$

  1. 逐位置计算:对于输出的每个位置 ,将所有通道对应窗口区域的元素与卷积核逐元素相乘并求和,得到该位置的输出值。

代码

import sys
input = sys.stdin.readline

def main():
    C, H_in, W_in = map(int, input().split())
    inp = []
    for c in range(C):
        channel = []
        for i in range(H_in):
            channel.append(list(map(int, input().split())))
        inp.append(channel)

    C2, K_h, K_w = map(int, input().split())
    kernel = []
    for c in range(C2):
        k_channel = []
        for i in range(K_h):
            k_channel.append(list(map(int, input().split())))
        kernel.append(k_channel)

    stride, padding = map(int, input().split())

    # 零填充
    H_pad = H_in + 2 * padding
    W_pad = W_in + 2 * padding
    padded = []
    for c in range(C):
        channel = [[0] * W_pad for _ in range(H_pad)]
        for i in range(H_in):
            for j in range(W_in):
                channel[i + padding][j + padding] = inp[c][i][j]
        padded.append(channel)

    H_out = (H_pad - K_h) // stride + 1
    W_out = (W_pad - K_w) // stride + 1

    result = [[0] * W_out for _ in range(H_out)]

    for i in range(H_out):
        for j in range(W_out):
            val = 0
            si = i * stride
            sj = j * stride
            for c in range(C):
                for ki in range(K_h):
                    for kj in range(K_w):
                        val += padded[c][si + ki][sj + kj] * kernel[c][ki][kj]
            result[i][j] = val

    for row in result:
        print(' '.join(map(str, row)))

main()

复杂度分析

  • 时间复杂度,对输出的每个位置,遍历所有通道和卷积核元素。
  • 空间复杂度,存储填充后的输入张量。