深入了解CNN及PyTorch中的层输出大小

卷积神经网络(CNN)在计算机视觉领域取得了巨大的成功,其核心在于对图像特征的提取和学习。理解CNN中每一层输出的大小,对于调试和优化模型至关重要。本文将讨论如何在PyTorch中输出每一层的大小,并提供代码示例以帮助读者更好地理解。

CNN基本要素

卷积神经网络主要由以下几层组成:

  1. 卷积层(Convolutional layer)
  2. 激活层(Activation layer)
  3. 池化层(Pooling layer)
  4. 全连接层(Fully connected layer)

每一层的输出大小依赖于输入大小、卷积核大小、步长(stride)、填充(padding)等因素。接下来,我们将用一个简单的CNN模型来演示如何在PyTorch中计算输出的大小。

PyTorch模型示例

以下是一个简单的CNN模型示例,展示了如何输出每一层的大小:

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(32 * 7 * 7, 120)
        self.fc2 = nn.Linear(120, 10)

    def forward(self, x):
        x = self.conv1(x)
        print(f'Conv1 output size: {x.size()}')  # 输出层大小
        x = F.relu(x)
        x = self.pool(x)
        print(f'After pooling output size: {x.size()}')  # 输出层大小
        x = self.conv2(x)
        print(f'Conv2 output size: {x.size()}')  # 输出层大小
        x = F.relu(x)
        x = self.pool(x)
        print(f'After second pooling output size: {x.size()}')  # 输出层大小
        x = x.view(-1, 32 * 7 * 7)  # 展平
        x = self.fc1(x)
        print(f'FC1 output size: {x.size()}')  # 输出层大小
        x = self.fc2(x)
        print(f'FC2 output size: {x.size()}')  # 输出层大小
        return x

# 测试模型
model = SimpleCNN()
input_tensor = torch.randn(1, 1, 28, 28)  # 输入尺寸为(1, 1, 28, 28)
output = model(input_tensor)

在这个示例中,我们创建了一个简单的CNN,包括两个卷积层和两个全连接层。我们在每层后面打印输出的大小,以便跟踪尺寸的变化。

每层输出尺寸的计算

我们通过卷积层和池化层来逐步计算输出尺寸。以第一层卷积层为例,输入为 28x28 的图像:

  • 卷积层输出尺寸计算: [ \text{Output Height} = \left\lfloor \frac{\text{Input Height} + 2 \times \text{Padding} - \text{Kernel Size}}{\text{Stride}} \right\rfloor + 1 ]

    conv1 为例: [ \text{Output Height} = \left\lfloor \frac{28 + 2 \times 1 - 3}{1} \right\rfloor + 1 = 28 ]

    所以在输出后图像的尺寸保持为 28x28。

  • 池化层输出尺寸计算: [ \text{Output Height} = \left\lfloor \frac{\text{Input Height} - \text{Kernel Size}}{\text{Stride}} \right\rfloor + 1 ]

    对于池化层: [ \text{Output Height} = \left\lfloor \frac{28 - 2}{2} \right\rfloor + 1 = 14 ]

这就是如何计算每一层的输出尺寸。

多层输出数据表(示例)

输入尺寸 输出尺寸
输入 (1, 1, 28, 28) (1, 16, 28, 28)
卷积层1 (1, 16, 28, 28) (1, 16, 28, 28)
池化层1 (1, 16, 28, 28) (1, 16, 14, 14)
卷积层2 (1, 16, 14, 14) (1, 32, 14, 14)
池化层2 (1, 32, 14, 14) (1, 32, 7, 7)
全连接层1 (1, 32, 7, 7) (1, 120)
全连接层2 (1, 120) (1, 10)

结论

理解CNN中每一层的输出大小对于构建和优化模型是至关重要的,因为它不仅影响最终的模型性能,还关系到后续维度的对齐和设计。通过在PyTorch中对每一层的输出进行打印,我们能够清晰地看到每一层的输出大小,帮助我们进行进一步的修改和优化。希望本文提供的代码示例和解释能够帮助你在进行深度学习时更好地理解CNN的工作原理。