PyTorch卷积神经网络训练时形状不匹配报错求助

问题根源 #

你的模型全连接层的输入维度是基于240×360的输入图像计算得到的86400，但实际训练数据集的图像尺寸为60×90，经过卷积和池化后，特征图展平后的维度变成了157696，和全连接层预期的输入维度不匹配，最终触发RuntimeError。

解决方案 #

1. 统一所有输入图像的尺寸（最简单直接） #

在数据集预处理阶段，强制将所有图像resize到模型设计时的目标尺寸（240×360），确保训练和测试时的输入尺寸一致：

from torchvision import transforms

# 定义预处理管道
train_transform = transforms.Compose([
    transforms.Resize((240, 360)),  # 统一为模型预期的输入大小
    transforms.ToTensor(),
    # 可按需添加其他预处理步骤（如归一化）
])

# 应用到数据集
train_dataset = YourDataset(root="path/to/data", transform=train_transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)

这样不管原始数据集的图像尺寸是多少，都会被调整为240×360，卷积池化后的特征维度会和测试时一致，全连接层就能正常接收输入。

2. 动态计算全连接层输入维度（更灵活） #

如果需要支持不同的输入尺寸，不要硬编码全连接层的输入维度，而是在模型初始化时通过一个dummy tensor自动计算卷积后的特征维度：

import torch
import torch.nn as nn

class SceneClassificationCNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        # 定义卷积池化层
        self.conv_block = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # 根据你的模型需求添加更多卷积池化层
        )
        
        # 用实际训练数据的尺寸创建dummy输入，计算卷积输出的特征维度
        dummy_input = torch.randn(1, 3, 60, 90)  # 对应数据集的60×90尺寸
        with torch.no_grad():
            conv_out = self.conv_block(dummy_input)
        # 计算单样本的特征展平后的维度
        self.flatten_dim = conv_out.numel() // dummy_input.size(0)
        
        # 定义全连接层，使用动态计算的维度
        self.fc_block = nn.Sequential(
            nn.Linear(self.flatten_dim, 512),
            nn.ReLU(),
            nn.Linear(512, num_classes)
        )
    
    def forward(self, x):
        x = self.conv_block(x)
        x = x.view(x.size(0), -1)  # 展平特征图
        x = self.fc_block(x)
        return x

这种方法会自动适配输入图像的尺寸，避免因硬编码维度导致的不匹配问题。

3. 检查并修正数据集预处理管道 #

确认DataLoader使用的transform中是否存在错误的resize操作，比如不小心将图像缩放到了60×90。如果是这样，直接修改Resize的参数为(240, 360)即可。

调试建议 #

• 验证输入尺寸：在训练前打印DataLoader中样本的形状，确认输入是否符合预期：

for images, labels in train_loader:
    print("Batch shape:", images.shape)  # 预期输出类似 (32, 3, 240, 360)
    break

• 跟踪特征维度变化：在模型的forward方法中打印每一层的输出形状，帮助定位尺寸不匹配的具体环节：

def forward(self, x):
    print("Input shape:", x.shape)
    x = self.conv_block[0](x)
    print("After first conv:", x.shape)
    x = self.conv_block[1](x)
    x = self.conv_block[2](x)
    print("After first pool:", x.shape)
    # 依次打印后续层的形状
    x = x.view(x.size(0), -1)
    print("Flattened shape:", x.shape)
    x = self.fc_block(x)
    return x

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