如何确定MNIST CNN分类器中nn.Linear与nn.Conv2d的关键参数？

问题描述 #

我已经针对MNIST数据集实现了CNN模型，能够理解CNN各层参数与形状的计算方式，但希望了解分类器部分（尤其是nn.Linear()）中in_features和out_features的确定方法，以及nn.Conv2d中in_channels和out_channels的选择依据。

实现代码 #

import torch.nn as nn

class CNNclf(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d((2, 2), stride=2),
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d((2, 2), stride=3),
            nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d((2, 2), stride=2))
        self.clf = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64, 20, bias=True),
            nn.ReLU(),
            nn.Linear(20, 10, bias=True))

    def forward(self, x):
        x = self.net(x)
        x = self.clf(x)
        return x

一、nn.Linear参数的确定方法 #

1. in_features：匹配卷积输出的扁平化维度

in_features必须严格等于卷积池化模块输出张量扁平化后的特征数量，这是硬性要求，否则会出现维度不匹配的错误。

以你的代码为例，一步步计算卷积池化后的输出形状（MNIST输入为(batch_size, 1, 28, 28)）：

• 第一层Conv2d（无padding、stride=1）：输出(batch_size, 64, 26, 26)（计算：28 - 3 + 1 = 26）
• 第一层MaxPool2d(2,2)：输出(batch_size, 64, 13, 13)（计算：26 / 2 = 13）
• 第二层Conv2d：输出(batch_size, 128, 11, 11)（计算：13 - 3 + 1 = 11）
• 第二层MaxPool2d(2,3)：输出(batch_size, 128, 4, 4)（计算：floor((11 - 2)/3) + 1 = 4）
• 第三层Conv2d：输出(batch_size, 64, 2, 2)（计算：4 - 3 + 1 = 2）
• 第三层MaxPool2d(2,2)：输出(batch_size, 64, 1, 1)（计算：2 / 2 = 1）

经过nn.Flatten()后，每个样本的特征维度为64 * 1 * 1 = 64，所以第一个nn.Linear的in_features设为64。

2. out_features：根据任务需求与模型容量调整

• 最后一层nn.Linear：out_features必须等于分类任务的类别总数，MNIST是10类手写数字，所以设为10，这是固定的。
• 中间隐藏层nn.Linear：属于超参数，通常在输入特征数和最终类别数之间做过渡（比如你这里从64降到20再到10）。选择时要平衡表达能力和过拟合风险：数值太小会导致特征提取不足，太大则会增加模型参数、减慢训练速度，还容易过拟合。简单任务用20-100之间的数值即可，也可以参考同类型任务的常用设置。

二、nn.Conv2d参数的选择依据 #

1. in_channels：完全由上一层输出决定

in_channels是固定值，没有选择空间：

• 第一层Conv2d：等于输入数据的通道数，MNIST是单通道灰度图，所以设为1；如果是RGB图则设为3。
• 后续Conv2d：必须等于前一层Conv2d的out_channels（池化层不改变通道数，所以如果前一层是池化，仍沿用之前Conv2d的out_channels）。比如你代码中第二层Conv2d的in_channels=64，对应第一层Conv2d的out_channels=64，这是必须匹配的，否则会报错。

2. out_channels：控制模型特征提取能力的超参数

out_channels决定了卷积层能提取的特征数量，选择时参考以下原则：

• 特征提取规律：通常遵循“先增后减”的模式——前期用更多通道提取丰富的底层、中层特征，后期减少通道精简核心特征（比如你代码中64→128→64的设置）。
• 任务复杂度匹配：简单任务（如MNIST）用较小的通道数（64、128）即可；复杂图像任务（如目标检测、图像分类）可能需要256、512甚至更高的通道数。
• 避免过拟合：通道数越多，模型参数越多，过拟合风险越高。如果训练数据量小，不要设置过大的通道数。
• 参考经典模型：比如LeNet用6、16通道，VGG系列用64、128、256、512的递增通道数，这些成熟设置可以直接参考。

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