PyTorch深度卷积神经网络维度不匹配问题调试求助

问题背景 #

要构建的CNN结构为：Conv2d => 激活层 => Conv2d => 激活层 => 最大池化 => Conv2d => 激活层 => Conv2d => 激活层 => 最大池化 => 全连接隐藏层 => 输出层。但创建modelDeeper = CNN(deeper=True).to(device)并训练时，出现以下错误：

RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x2048 and 128x512).

原CNN类代码如下：

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

class CNN(nn.Module):

    def __init__(self,dropout = 0, silu = False, numHidden = 1, deeper = False):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 5, padding='same')
        self.conv2 = nn.Conv2d(16, 32, 5, padding='same')
        self.convDeeper1 = nn.Conv2d(3, 32, 3, padding='same')
        self.convDeeper2 = nn.Conv2d(32, 32, 3, padding='same')

        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32*8*8,512)
        self.fcDeeper1 = nn.Linear(128, 512) # not sure why the input dimension should be 128
        
        self.fc2 = nn.Linear(512, 10)
        self.dropoutMaxPool = nn.Dropout(0)
        self.dropoutFC = nn.Dropout(0)
        if dropout != 0:
          self.dropoutMaxPool = nn.Dropout(0.25)
          self.dropoutFC = nn.Dropout(0.5)
        self.silu = silu
        self.numHidden = numHidden
        self.deeper = deeper

    def activation(self, x):
      if not self.silu:
        return F.relu(x)
      else:
        return F.silu(x)

    def forward(self, x):
        if not self.deeper:
          x = self.activation(self.conv1(x))
          x = self.pool(x) # max pool layer
          x = self.dropoutMaxPool(x) # insert 

          x = self.activation(self.conv2(x))
          x = self.pool(x) 
          x = self.dropoutMaxPool(x)

        else:
          x = self.activation(self.convDeeper1(x))
          x = self.pool(x) # max pool layer
          x = self.dropoutMaxPool(x) # insert 

          x = self.activation(self.convDeeper2(x))
          x = self.pool(x) 
          x = self.dropoutMaxPool(x)

          x = self.activation(self.convDeeper2(x))
          x = self.pool(x) # max pool layer
          x = self.dropoutMaxPool(x) # insert 

          x = self.activation(self.convDeeper2(x))
          x = self.pool(x) 
          x = self.dropoutMaxPool(x)

        x = torch.flatten(x, 1) # flatten all dimensions except batch
        
        x = self.activation(self.fcDeeper1(x)) # does this use rectified linear units? This is the hidden fully connected layer.
        x = self.dropoutFC(x)
        
        

        x = F.log_softmax(self.fc2(x), 1) # for the output
        return x

错误原因 #

1. 结构不符：deeper=True分支的网络做了4次池化，和预期的2次池化结构不一致，导致特征图尺寸计算错误。
2. 维度不匹配：全连接层fcDeeper1的输入维度设为128，但实际卷积+池化后展平的特征数是2048，矩阵乘法时维度无法对应。

修正方案 #

1. 对齐网络结构到预期 #

将deeper=True分支改为两组Conv→激活→Conv→激活→池化的结构，匹配需求：

else:
    # 第一组：Conv→激活→Conv→激活→池化
    x = self.activation(self.convDeeper1(x))
    x = self.activation(self.convDeeper2(x))
    x = self.pool(x)
    x = self.dropoutMaxPool(x)

    # 第二组：Conv→激活→Conv→激活→池化
    x = self.activation(self.convDeeper2(x))
    x = self.activation(self.convDeeper2(x))
    x = self.pool(x)
    x = self.dropoutMaxPool(x)

2. 修正全连接层输入维度 #

假设输入是3通道32x32图像（如CIFAR10），两次池化后特征图尺寸为8x8，通道数32，展平后特征数为32*8*8=2048，修改fcDeeper1的初始化：

self.fcDeeper1 = nn.Linear(32*8*8, 512)  # 替换原128为2048

完整修正后的代码 #

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

class CNN(nn.Module):

    def __init__(self,dropout = 0, silu = False, numHidden = 1, deeper = False):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 5, padding='same')
        self.conv2 = nn.Conv2d(16, 32, 5, padding='same')
        self.convDeeper1 = nn.Conv2d(3, 32, 3, padding='same')
        self.convDeeper2 = nn.Conv2d(32, 32, 3, padding='same')

        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32*8*8,512)
        self.fcDeeper1 = nn.Linear(32*8*8, 512) 
        
        self.fc2 = nn.Linear(512, 10)
        self.dropoutMaxPool = nn.Dropout(0)
        self.dropoutFC = nn.Dropout(0)
        if dropout != 0:
          self.dropoutMaxPool = nn.Dropout(0.25)
          self.dropoutFC = nn.Dropout(0.5)
        self.silu = silu
        self.numHidden = numHidden
        self.deeper = deeper

    def activation(self, x):
      if not self.silu:
        return F.relu(x)
      else:
        return F.silu(x)

    def forward(self, x):
        if not self.deeper:
          x = self.activation(self.conv1(x))
          x = self.pool(x) 
          x = self.dropoutMaxPool(x)

          x = self.activation(self.conv2(x))
          x = self.pool(x) 
          x = self.dropoutMaxPool(x)

        else:
          # 第一组：Conv→激活→Conv→激活→池化
          x = self.activation(self.convDeeper1(x))
          x = self.activation(self.convDeeper2(x))
          x = self.pool(x) 
          x = self.dropoutMaxPool(x)

          # 第二组：Conv→激活→Conv→激活→池化
          x = self.activation(self.convDeeper2(x))
          x = self.activation(self.convDeeper2(x))
          x = self.pool(x) 
          x = self.dropoutMaxPool(x)

        x = torch.flatten(x, 1)
        
        x = self.activation(self.fcDeeper1(x))
        x = self.dropoutFC(x)
        
        x = F.log_softmax(self.fc2(x), 1)
        return x

额外验证方法 #

如果输入图像尺寸不是32x32，可在forward中添加打印代码确认展平后的维度：

x = torch.flatten(x, 1)
print(x.shape) # 输出格式为(batch_size, 特征数)，用特征数调整全连接层输入维度

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