PyTorch训练不同尺寸图像数据集报张量尺寸不匹配的解决方法

报错核心原因：PyTorch默认DataLoader的default_collate拼接函数会将同批次所有张量沿batch维度堆叠，强制要求同批次张量尺寸完全一致。你的数据集图像尺寸从[3, 300, 535]到[3, 1080, 1920]不等，直接堆叠必然触发尺寸不匹配的RuntimeError。
调整batch size、kernel size完全解决不了这个问题：batch size仅控制单批次样本数量，kernel size是卷积操作的窗口大小，二者都不能让不同尺寸的图像自动对齐。此外你当前的神经网络存在通道数不匹配、卷积核尺寸过大的硬伤，就算解决拼接问题也无法正常运行。

方案1：统一输入尺寸（最通用，90%以上视觉任务推荐使用） #

通过预处理变换把所有图像缩放到固定尺寸，是最简单、训练效率最高的方案：

• 第一步：给数据集传入标准化预处理变换，修改主文件里的数据集初始化代码

from torchvision import transforms

# 定义预处理流水线：统一缩放到224*224，转张量、归一化
base_transform = transforms.Compose([
    transforms.ToPILImage(),  # 适配read_image读取的uint8张量格式
    transforms.Resize((224, 224)),  # 所有图像强制对齐到224*224，可根据任务调整尺寸
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 初始化数据集时传入transform
test_dataset = dataset.csHeadBody(
    csv_file="images\\test_labels.csv",
    root_dir="images\\test",
    transform=base_transform
)
train_dataset = dataset.csHeadBody(
    csv_file="images\\train_labels.csv",
    root_dir="images\\train",
    transform=base_transform
)

• 第二步：修复神经网络结构错误，替换原有不合理的卷积、全连接层配置

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

class Net(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        # 224*224输入经过3次2*2池化后尺寸为28*28，通道数64
        self.fc1 = nn.Linear(64 * 28 * 28, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

• 补充：你的数据集定义文件缺少必要依赖导入，需要在文件开头加上以下代码，否则初始化时会报名称错误：

import os
import pandas as pd
from torchvision.io import read_image
from torch.utils.data import Dataset

方案2：自定义批次拼接逻辑+自适应池化（不想强制拉伸图像时使用） #

如果担心强制resize破坏图像宽高比，可以自定义collate_fn，配合网络中的自适应池化层实现变尺寸训练：

• 第一步：自定义collate函数，初始化DataLoader时传入

from torchvision import transforms

def variable_size_collate(batch):
    imgs = []
    labels = []
    for img, label in batch:
        # 短边对齐到256像素，保持原始宽高比
        _, h, w = img.shape
        scale = 256 / min(h, w)
        new_h = int(h * scale)
        new_w = int(w * scale)
        img = transforms.functional.resize(img, (new_h, new_w))
        imgs.append(img)
        labels.append(label)
    labels = torch.tensor(labels, dtype=torch.long)
    return imgs, labels

# 初始化DataLoader时指定自定义拼接函数
train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=batch_size,
    shuffle=True,
    collate_fn=variable_size_collate
)
test_loader = DataLoader(
    dataset=test_dataset,
    batch_size=batch_size,
    shuffle=True,
    collate_fn=variable_size_collate
)

• 第二步：修改网络结构，加入自适应池化层统一特征图尺寸

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

class Net(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        # 自适应池化：无论输入特征图尺寸多大，统一输出7*7大小
        self.adaptive_pool = nn.AdaptiveAvgPool2d((7,7))
        self.fc1 = nn.Linear(64 * 7 * 7, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

    def forward(self, x):
        batch_features = []
        # 逐张提取卷积特征
        for img in x:
            img = img.unsqueeze(0) # 补充单样本batch维度
            feat = self.pool(F.relu(self.conv1(img)))
            feat = self.pool(F.relu(self.conv2(feat)))
            feat = self.pool(F.relu(self.conv3(feat)))
            feat = self.adaptive_pool(feat)
            batch_features.append(feat)
        # 所有特征对齐后再拼接成batch
        x = torch.cat(batch_features, dim=0)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

注意：方案2训练速度比方案1慢30%以上，因为需要逐张做卷积前向，没有批量计算的优势，非必要优先选方案1。

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