PyTorch训练自定义数据集报错：矩阵形状不匹配求助

问题描述 #

我尝试用PyTorch训练自定义图像分类模型，数据集包含train_data和test_data两个文件夹，每个文件夹下有detected和not_detected两类子文件夹，内部是不同尺寸的PNG图片。之前用MNIST数据集时代码正常运行，但使用自定义数据集时，即便通过transforms.Resize((256,256))和RandomCrop((224,224))处理图像，仍出现如下错误：

RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x150528 and 784x512)

原代码 #

import torch
import torch.nn as nn
from torch.utils.data import Dataset
from torchvision import datasets
import torch.optim as optim
import os
from PIL import Image

import torchvision.transforms as transforms
from torch.utils.data import DataLoader

is_gpu = torch.cuda.is_available()

class CustomDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.classes = os.listdir(root_dir)
        self.data = []
        self.targets = []

        for class_idx, class_name in enumerate(self.classes):
            class_dir = os.path.join(root_dir, class_name)
            for filename in os.listdir(class_dir):
                img_path = os.path.join(class_dir, filename)
                self.data.append(img_path)
                self.targets.append(class_idx)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path = self.data[idx]
        target = self.targets[idx]

        img = Image.open(img_path).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        return img, target

transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.RandomCrop((224, 224)),
    transforms.ToTensor()
])

train_dataset = CustomDataset(root_dir=r'C:\Users\enis_\Desktop\yololu\train_data', transform=transform)
test_dataset = CustomDataset(root_dir=r'C:\Users\enis_\Desktop\yololu\test_data', transform=transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)


model = nn.Sequential(nn.Flatten(),
                     nn.Linear(28*28,512),
                     nn.ReLU(),
                     nn.Linear(512,256),
                     nn.ReLU(),
                     nn.Linear(256,10))

if is_gpu:
    model = model.cuda()

print(model)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),lr=0.7)

epochs = 5

for i in range(epochs):
    train_loss = 0.0
    for data,label in train_loader:
        if is_gpu:
            data, label = data.cuda(), label.cuda()
        optimizer.zero_grad()
        
        output = model(data)
        loss = criterion(output,label)
        loss.backward()
        
        optimizer.step()
        
        train_loss += loss.item() * data.size(0)
    print(f'Epoch: {i+1} / {epochs} \t\t\t Training Loss:{train_loss/len(train_loader)}')

完整报错信息 #

Sequential(
(0): Flatten(start_dim=1, end_dim=-1)
(1): Linear(in_features=784, out_features=512, bias=True)
(2): ReLU()
(3): Linear(in_features=512, out_features=256, bias=True)
(4): ReLU()
(5): Linear(in_features=256, out_features=10, bias=True)
)
Traceback (most recent call last):
File "c:\Users\enis_\Desktop\yololu\tempCodeRunnerFile.py", line 80, in 
output = model(data)
^^^^^^^^^^^
File "C:\Users\enis_\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\enis_\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\torch\nn\modules\container.py", line 217, in forward
input = module(input)
^^^^^^^^^^^^^
File "C:\Users\enis_\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\enis_\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\torch\nn\modules\linear.py", line 114, in forward
return F.linear(input, self.weight, self.bias)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x150528 and 784x512)

错误原因分析 #

• 原模型是为MNIST单通道28x28图像设计的，第一层Linear输入为28*28=784，但自定义数据是RGB三通道，经过裁剪后尺寸为224x224，展平后特征数为3*224*224=150528，与Linear层输入维度不匹配，导致矩阵乘法错误。
• 最后一层Linear输出设置为10类，但数据集只有2类，后续损失计算也会出现逻辑错误。
• SGD优化器学习率设置为0.7过高，会导致模型训练过程中损失剧烈震荡，难以收敛。

修改后的代码 #

import torch
import torch.nn as nn
from torch.utils.data import Dataset
import torch.optim as optim
import os
from PIL import Image
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

is_gpu = torch.cuda.is_available()

class CustomDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.classes = os.listdir(root_dir)
        self.data = []
        self.targets = []

        for class_idx, class_name in enumerate(self.classes):
            class_dir = os.path.join(root_dir, class_name)
            for filename in os.listdir(class_dir):
                img_path = os.path.join(class_dir, filename)
                self.data.append(img_path)
                self.targets.append(class_idx)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path = self.data[idx]
        target = self.targets[idx]

        img = Image.open(img_path).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        return img, target

transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.RandomCrop((224, 224)),
    transforms.ToTensor()
])

train_dataset = CustomDataset(root_dir=r'C:\Users\enis_\Desktop\yololu\train_data', transform=transform)
test_dataset = CustomDataset(root_dir=r'C:\Users\enis_\Desktop\yololu\test_data', transform=transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 适配RGB 224x224图像与二分类任务的模型
model = nn.Sequential(nn.Flatten(),
                     nn.Linear(3*224*224, 512),  # 输入特征数匹配RGB图像展平后维度
                     nn.ReLU(),
                     nn.Linear(512, 256),
                     nn.ReLU(),
                     nn.Linear(256, 2))  # 输出类别数匹配数据集的两类

if is_gpu:
    model = model.cuda()

print(model)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)  # 调整学习率到合理范围

epochs = 5

for i in range(epochs):
    train_loss = 0.0
    for data,label in train_loader:
        if is_gpu:
            data, label = data.cuda(), label.cuda()
        optimizer.zero_grad()
        
        output = model(data)
        loss = criterion(output,label)
        loss.backward()
        
        optimizer.step()
        
        train_loss += loss.item() * data.size(0)
    # 修正损失计算方式：除以数据集总长度得到平均样本损失
    print(f'Epoch: {i+1} / {epochs} \t\t\t Training Loss:{train_loss/len(train_loader.dataset):.4f}')

关键修改点 #

• 第一层Linear输入特征数改为3*224*224，匹配RGB图像展平后的维度
• 最后一层Linear输出改为2，对应数据集的二分类任务
• 将SGD学习率从0.7调整为0.01，避免训练过程中损失震荡
• 修正训练损失计算方式：除以数据集总长度，结果更直观反映平均样本损失

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