如何解决卷积层Flatten后到前馈层的输入维度不匹配问题？

核心问题 #

你遇到的维度不匹配问题确实是LayerChoice选择不同卷积核导致的：

• CIFAR10输入是32×32的图片
• 选3×3卷积（padding=1，stride=1）：输出尺寸仍为32×32
• 选5×5卷积（padding=1，stride=1）：输出尺寸变为30×30（计算公式：(32 -5 +2×1)/1 +1 =30）
  后续经过conv2、max_pool2d后，两种路径的特征图维度完全不同，导致flatten后的维度不一致，和硬编码的14400无法匹配。

解决方法 #

方法1：用自适应池化固定特征图维度 #

在flatten前加入自适应池化层，将特征图统一为固定尺寸，这样不管前面卷积路径如何选择，flatten后的维度都是固定值。修改后的模型代码：

from nni.nas.pytorch import mutables
import torch.nn as nn
import torch.nn.functional as F
import torch

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = mutables.LayerChoice([
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(3, 32, kernel_size=5, stride=1, padding=1)
        ])
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        # 自适应池化将特征图固定为15×15，flatten后为64×15×15=14400
        self.adaptive_pool = nn.AdaptiveAvgPool2d((15,15))
        self.fc1 = nn.Linear(14400, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = self.adaptive_pool(x)  # 新增自适应池化
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

方法2：用懒初始化层自动推断输入维度 #

使用PyTorch提供的nn.LazyLinear，仅需指定输出单元数，输入维度会在第一次前向传播时自动推断，完美适配NAS场景中输入维度变化的情况：

from nni.nas.pytorch import mutables
import torch.nn as nn
import torch.nn.functional as F
import torch

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = mutables.LayerChoice([
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(3, 32, kernel_size=5, stride=1, padding=1)
        ])
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        # LazyLinear仅需指定输出维度，输入维度自动推断
        self.fc1 = nn.LazyLinear(128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

注意：nn.LazyLinear需要PyTorch 1.8及以上版本支持，第一次前向传播完成后权重会固定，后续不可再改变输入维度。

关于PyTorch是否支持类似Keras的全连接层设置 #

是的，PyTorch从1.8版本开始提供了nn.LazyLinear（以及nn.LazyConv2d等系列懒初始化层），用法和Keras一致，仅需指定输出单元数，输入维度会在第一次前向传播时自动推断。但要注意，懒初始化层在完成第一次前向传播前没有权重，无法直接调用parameters()或保存模型。

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