MNIST数据集CNN中首个全连接层输入维度如何正确计算？

全连接层输入维度计算方法 #

卷积、池化层输出的张量维度格式为[batch_size, 通道数, 特征图高度, 特征图宽度]，你打印的torch.Size([10,16,4,4])四个维度的含义分别是：

• 10：当前批次的样本数量（batch size）
• 16：最后一层卷积输出的通道数
• 4：输出特征图的高度
• 4：输出特征图的宽度

首个全连接层要求输入单样本的一维特征向量，因此单样本的输入维度为 通道数 × 特征图高度 × 特征图宽度，对应你的场景就是 16 × 4 × 4 = 256，你代码里写的in_features=16*4*4这个数值本身是正确的，报错的核心原因是你没有把卷积输出的四维张量展平为二维张量，就直接传入了全连接层。

修复方案 #

你需要在前向传播逻辑中，放开X = self.fc1(X)之前注释掉的X = X.view()代码，修改为以下写法完成张量展平：

# 写法1：自动适配batch大小，将后面三个维度合并为一维特征
X = X.view(X.size(0), -1)
# 写法2：明确指定展平后的特征维度，适配当前输出尺寸
X = X.view(-1, 16*4*4)

修改后的完整前向传播代码示例：

def forward_pass(self, X):
    # 第1层卷积
    X = self.conv1(X)
    X = self.reLU1(X)
    X = self.MaxPool1(X)
    # 第2层卷积
    X = self.conv2(X)
    X = self.reLU2(X)
    X = self.MaxPool2(X)

    # 新增展平操作
    X = X.view(X.size(0), -1)
    X = self.fc1(X)
    X = self.sig(X)
    X = self.fc2(X)
  
    return X

内容的提问来源于stack exchange，提问作者Andy321