PyTorch CNN修改卷积参数后出现输入尺寸无效错误求助

问题背景 #

原PyTorch CNN代码可正常运行，修改卷积层参数后出现形状不匹配错误：

ERROR in CNN Pytorch; shape '[-1, 192]' is invalid for input of size 300000

修改的代码片段：

self.conv1 = nn.Conv2d(in_channels=1,out_channels=8,kernel_size=3)
self.conv2 = nn.Conv2d(in_channels=8,out_channels=16,kernel_size=3)
self.fc1 = nn.Linear(in_features=16*2*2,out_features=128)

错误原因 #

1. 特征图尺寸计算错误：FashionMNIST输入为28×28单通道图像，修改卷积核后特征图尺寸计算偏差，导致全连接层输入特征数设置错误。
2. forward函数reshape参数未同步更新：原代码中x.reshape(-1,12*4*4)未随网络结构修改，导致张量展平后的维度与fc1输入维度不匹配。

正确尺寸计算 #

• 输入图像：28×28
• conv1（kernel=3, padding=0, stride=1）输出尺寸：28 - 3 + 1 = 26 → 经maxpool2d(2,2)后变为13×13
• conv2（kernel=3, padding=0, stride=1）输出尺寸：13 - 3 + 1 = 11 → 经maxpool2d(2,2)后变为5×5（PyTorch默认向下取整）
• 因此fc1输入特征数应为：16（通道数）×5×5 = 400

修正后的完整代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import FashionMNIST
from torchvision import transforms
import torch.optim as optim

class Network(nn.Module):
  def __init__(self):
    super(Network,self).__init__()
    self.conv1 = nn.Conv2d(in_channels=1,out_channels=8,kernel_size=3)
    self.conv2 = nn.Conv2d(in_channels=8,out_channels=16,kernel_size=3)
    # 修正fc1输入特征数为16*5*5=400
    self.fc1 = nn.Linear(in_features=16*5*5,out_features=128)
    self.fc2 = nn.Linear(in_features=128,out_features=64)
    self.out = nn.Linear(in_features=64,out_features=10)
    
  def forward(self,x):
    # input layer
    x = x
    # first hidden layer
    x = self.conv1(x)
    x = F.relu(x)
    x = F.max_pool2d(x,kernel_size=2,stride=2)
    # second hidden layer
    x = self.conv2(x)
    x = F.relu(x)
    x = F.max_pool2d(x,kernel_size=2,stride=2)
    # third hidden layer
    # 修正reshape的特征数为400
    x = x.reshape(-1,16*5*5)
    x = self.fc1(x)
    x = F.relu(x)
    # fourth hidden layer
    x = self.fc2(x)
    x = F.relu(x)
    
    # output layer
    x = self.out(x)
    return x


batch_size = 1000
train_dataset = FashionMNIST(
    '../data', train=True, download=True, 
    transform=transforms.ToTensor())
trainloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

test_dataset = FashionMNIST(
    '../data', train=False, download=True, 
    transform=transforms.ToTensor())
testloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

model = Network()

losses = []
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())
epochs = 1

for i in range(epochs):
    batch_loss = []
    for j, (data, targets) in enumerate(trainloader):
        optimizer.zero_grad()
        ypred = model(data)
        loss = criterion(ypred, targets.reshape(-1))
        loss.backward()
        optimizer.step()
        batch_loss.append(loss.item())
    if i>10: 
        optimizer.param_groups[0]['lr'] = 0.0005  # 修正原代码中学习率修改方式的错误
    losses.append(sum(batch_loss) / len(batch_loss))
    print('Epoch {}: loss {:.4f}'.format(i, losses[-1]))

额外说明 #

• 原代码中optimizer.lr = 0.0005的学习率修改方式错误，PyTorch需通过optimizer.param_groups[0]['lr']调整
• 若不确定特征图尺寸，可在forward函数中添加print(x.shape)打印各层输出形状，方便调试

内容的提问来源于stack exchange，提问作者saif ul islam