PyTorch中送入Linear层前如何计算输出图像形状？

这个5是经过两次卷积+池化后，单通道特征图的边长，计算基于卷积和池化层的输出尺寸规则，我们以CIFAR-10数据集的32×32×3输入图片为例（这是该网络的典型输入场景），一步步推导：

卷积层输出尺寸计算公式 #

代码中nn.Conv2d未指定padding和stride时，默认参数为padding=0、stride=1，输出特征图边长公式为：

输出边长 = (输入边长 - 卷积核边长 + 2*padding) // stride + 1

池化层输出尺寸计算公式 #

代码中nn.MaxPool2d(2, 2)表示池化核边长为2、步长为2，输出边长公式为：

输出边长 = 输入边长 // 池化核边长

逐层计算过程 #

• 第一步：输入→conv1→pool
  输入图片边长：32
  conv1（核边长5）输出边长：(32 - 5 + 0) // 1 + 1 = 28
  经过pool（核边长2）后边长：28 // 2 = 14
  此时特征图尺寸：14×14×6（6是conv1的输出通道数）

• 第二步：pool输出→conv2→pool
  conv2（核边长5）输出边长：(14 - 5 + 0) // 1 + 1 = 10
  再次经过pool后边长：10 // 2 = 5
  此时特征图尺寸：5×5×16（16是conv2的输出通道数）

全连接层输入维度 #

当把特征图展平为一维向量时，总维度就是通道数 × 单通道边长 × 单通道边长，也就是16 * 5 * 5，这就是fc1输入维度的由来。

内容的提问来源于stack exchange，提问作者Arup Sankar Roy