PyTorch实现多输入输出通道的卷积操作

news2025/4/19 3:05:19

本文通过代码示例详细讲解如何在PyTorch中实现多输入通道和多输出通道的卷积运算，并对比传统卷积与1x1卷积的实现差异。

1. 多输入通道互相关运算

当输入包含多个通道时，卷积核需要对每个通道分别进行互相关运算，最后将结果相加。以下是实现代码：

import torch
from d2l import torch as d2l

def corr2d_multi_in(X, K):
    return sum(d2l.corr2d(x, k) for x, k in zip(X, K))

验证输出：
输入一个2通道的3x3张量和一个2通道的2x2卷积核，输出结果为2x2张量：

X = torch.tensor([[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]],
                 [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]])
K = torch.tensor([[[0.0, 1.0], [2.0, 3.0]], [[1.0, 2.0], [3.0, 4.0]]])

print(corr2d_multi_in(X, K))

输出结果：

tensor([[ 56.,  72.],
        [104., 120.]])

2. 多输出通道互相关运算

通过堆叠多个卷积核，可以实现多输出通道。以下代码展示了如何生成3个输出通道：

def corr2d_multi_in_out(X, K):
    return torch.stack([corr2d_multi_in(X, k) for k in K], 0)

K = torch.stack((K, K+1, K+2), 0)  # 堆叠3个卷积核
print("卷积核形状:", K.shape)

输出结果：

卷积核形状: torch.Size([3, 2, 2, 2])

运行多通道卷积：

print(corr2d_multi_in_out(X, K))

输出结果：

tensor([[[ 56.,  72.],
         [104., 120.]],

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

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

3. 1x1卷积的优化实现

1x1卷积可通过矩阵乘法高效实现，尤其适用于通道维度调整。以下是对比传统卷积与1x1卷积的代码：

def corr2d_multi_in_out_1x1(X, K):
    c_i, h, w = X.shape
    c_o = K.shape[0]
    X = X.reshape((c_i, h * w))       # 展平空间维度
    K = K.reshape((c_o, c_i))        # 展平卷积核
    Y = torch.matmul(K, X)           # 矩阵乘法
    return Y.reshape((c_o, h, w))    # 恢复形状

# 生成随机输入和卷积核
X = torch.normal(0, 1, (3, 3, 3))    # 3通道3x3输入
K = torch.normal(0, 1, (2, 3, 1, 1)) # 2输出通道的1x1卷积核

# 验证两种方法结果一致
Y1 = corr2d_multi_in_out_1x1(X, K)
Y2 = corr2d_multi_in_out(X, K)
assert float(torch.abs(Y1 - Y2).sum()) < 1e-6  # 误差极小