PyTorch中Conv2D层批量维度指定与矩阵维度不匹配报错排查

核心错误原因 #

• 缺少特征展平操作：卷积和池化后输出的是4维张量[batch_size, channels, height, width]，直接接入全连接层时，PyTorch的nn.Linear默认对最后一个维度做线性变换，而非合并通道、高、宽维度为单一特征维度。这导致输入到Linear层的张量被错误解析为[50*16*100, 100]（即80000x100），与你定义的Linear层输入维度160000不匹配，触发维度错误。
• 损失函数输入输出维度不匹配：即使解决了前序维度问题，后续MSELoss也会报错——模型输出是[50,1000]，而标签x是[50,1,600,600]，两者形状完全不一致。

修正后的代码 #

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import FakeData
from torchvision.transforms import ToTensor

# 生成模拟数据集
dataset = FakeData(size=500, image_size=(1, 600, 600), transform=ToTensor())
training_data, test_data = random_split(dataset, [400, 100])
train_dataloader = DataLoader(training_data, batch_size=50, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=50, shuffle=True)

net = nn.Sequential(
    nn.Conv2d(
        in_channels=1,              
        out_channels=16,            
        kernel_size=5,                     
        padding=2,           
    ),
    nn.ReLU(),  
    nn.MaxPool2d(kernel_size=6),
    nn.Flatten(),  # 新增：将4维特征张量展平为2维[batch_size, channels*height*width]
    nn.Linear(160000, 1000),
    nn.ReLU(),
    # 图像重建任务需添加：将特征映射回输入图像维度
    nn.Linear(1000, 1*600*600),
    nn.Unflatten(1, (1, 600, 600))  # 将展平特征恢复为图像形状
)

optimizer = optim.Adam(net.parameters(), lr=1e-3,)
criterion = nn.MSELoss()  # 提前定义损失函数，避免重复创建

epochs = 10
for i in range(epochs):
    for (x, _) in train_dataloader:
        optimizer.zero_grad()

        print(x.shape) # 输出torch.Size([50, 1, 600, 600])
        output = net(x)

        loss = criterion(output, x)
        loss.backward()
        optimizer.step()
    print(f"Epoch {i+1}, Loss: {loss.item()}")

关键修正说明 #

• 添加nn.Flatten()：在池化层与全连接层之间插入该层，将[50,16,100,100]转换为[50,160000]，匹配全连接层的输入要求。
• 修复损失函数维度：如果是图像重建任务，需在最后添加线性层将1000维特征映射回1*600*600，再用nn.Unflatten恢复为图像形状，确保与输入x维度一致。
• 优化代码结构：将损失函数定义移到循环外，避免每次迭代重复创建实例，提升运行效率。

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