卷积层的多输入与多输出通道

import torch
from d2l import torch as d2l

# 输入为多个通道，输出为一个通道，将多个通道分别进行卷积再相加
def corr2d_multi_in(X, K):
    # 先遍历“X”和“K”的第0个维度（通道维度），再把它们加在一起
    return sum(d2l.corr2d(x, k) for x, k in zip(X, K))
# X的尺寸为（2，3，3）
X = torch.tensor([[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]],      #通道1
                  [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]])     #通道2
# K的尺寸为（2，2，2）
K = torch.tensor([[[0.0, 1.0], [2.0, 3.0]],     #通道1
                  [[1.0, 2.0], [3.0, 4.0]]])    #通道2
# Y的尺寸为（2，2）
Y = corr2d_multi_in(X, K)
print(Y)
# 输出：
# tensor([[ 56.,  72.],
#         [104., 120.]])

# --------------------------------------------------------------------------------
# 输入为多个通道，输出为多个通道，将多个通道分别与某一个卷积核进行卷积再相加得到一个输出通道，并将多个输出通道拼接起来
def corr2d_multi_in_out(X, K):
    # 迭代“K”的第0个维度，每次都对输入“X”执行互相关运算。
    # 最后将所有结果都叠加在一起
    return torch.stack([corr2d_multi_in(X, k) for k in K], 0)
# 将原卷积核K、原卷积核每个元素加1得到的新卷积核K+1、原卷积核每个元素加2得到的新卷积核K+2进行拼接得到一个具有三个输出通道的卷积核
K = torch.stack((K, K + 1, K + 2), 0)
Y = corr2d_multi_in_out(X, K)
print(Y)
# 输出：
# tensor([[[ 56.,  72.],
#          [104., 120.]],

#         [[ 76., 100.],
#          [148., 172.]],

#         [[ 96., 128.],
#          [192., 224.]]])