PyTorch如何支持语义分割任务中不同尺寸的模型输入数据集？

该报错和你的模型是否支持任意尺寸输入无关，本质是**PyTorch DataLoader的默认拼接逻辑（collate_fn）**需要同一个批次内的所有图像、掩码张量尺寸完全一致，才能通过torch.stack操作新增batch维度拼接为批量张量，你的数据集内图像尺寸不一致，在数据加载阶段就触发了报错，还没到模型推理环节。

你可以根据自己的需求选择以下任意一种方案：

方案1：自定义collate_fn返回样本列表（最适配任意尺寸模型） #

无需修改原始图像尺寸，只需自定义拼接逻辑，不强制堆叠为批量张量：

1. 先实现自定义拼接函数：

def custom_collate(batch):
    # 入参batch是列表，每个元素为数据集返回的(图像张量, 掩码张量)
    images = [item[0] for item in batch]
    masks = [item[1] for item in batch]
    return images, masks

2. 在DataLoader中指定该自定义函数：

train_dataloader = DataLoader(train_dataset,batch_size=BATCH_SIZE,shuffle=True, collate_fn=custom_collate)
val_dataloader = DataLoader(val_dataset,batch_size=BATCH_SIZE,shuffle=True, collate_fn=custom_collate)

3. 修改训练/验证循环的推理逻辑，遍历批次内的每个样本单独推理：

def train_fn(dataloader, model, loss_fn, optimizer, scaler):
    size = len(dataloader.dataset)
    loop = tqdm(dataloader)
    for batch,(data,targets) in enumerate(loop):
        total_loss = 0
        # 逐个样本推理
        for img, mask in zip(data, targets):
            img = img.to(device=DEVICE).unsqueeze(0) # 手动补充batch维度
            mask = mask.to(device=DEVICE).unsqueeze(0)
            with torch.cuda.amp.autocast():
                predictions = model(img)
                loss = loss_fn(predictions,targets)
            total_loss += loss
        # 计算平均损失再反向传播
        avg_loss = total_loss / len(data)
        optimizer.zero_grad()
        if DEVICE == "cuda":
            scaler.scale(avg_loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            avg_loss.backward()
            optimizer.step()
        loop.set_postfix(loss=avg_loss.item())

验证函数也需要做对应修改。

方案2：动态padding到批次最大尺寸（兼顾批量推理效率） #

既保留批量推理的效率，也不用修改原始图像的长宽比例，只需将每个批次内的图像、掩码padding到当前批次的最大尺寸：

1. 实现带padding的自定义collate_fn：

def pad_collate(batch, img_pad_val=0, mask_pad_val=255):
    images = [item[0] for item in batch]
    masks = [item[1] for item in batch]
    # 取当前批次的最大高、宽
    max_h = max([img.shape[1] for img in images])
    max_w = max([img.shape[2] for img in images])
    padded_imgs, padded_masks = [], []
    for img, mask in zip(images, masks):
        h, w = img.shape[1], img.shape[2]
        pad_h, pad_w = max_h - h, max_w - w
        # 图像padding补0
        padded_img = torch.nn.functional.pad(img, (0, pad_w, 0, pad_h), value=img_pad_val)
        padded_imgs.append(padded_img)
        # 掩码padding补损失函数的忽略值，避免padding区域影响训练
        padded_mask = torch.nn.functional.pad(mask, (0, pad_w, 0, pad_h), value=mask_pad_val)
        padded_masks.append(padded_mask)
    # 尺寸统一后再堆叠为批量张量
    return torch.stack(padded_imgs), torch.stack(padded_masks)

2. 在损失函数中指定忽略值，在DataLoader中指定该collate_fn：

# 损失函数新增ignore_index参数，和mask_pad_val保持一致
loss_fn = nn.CrossEntropyLoss(weight = class_weights, ignore_index=255)
# DataLoader指定自定义collate_fn
train_dataloader = DataLoader(train_dataset,batch_size=BATCH_SIZE,shuffle=True, collate_fn=pad_collate)
val_dataloader = DataLoader(val_dataset,batch_size=BATCH_SIZE,shuffle=True, collate_fn=pad_collate)

该方案无需修改训练/验证的循环逻辑，返回的仍是正常的4维批量张量。

方案3：统一图像尺寸（最简单通用） #

如果可以接受修改原始图像尺寸，直接在数据预处理阶段加Resize/Crop操作，将所有图像统一为固定尺寸即可：

from torchvision import transforms
# 预处理逻辑：统一缩放到512*512，也可根据需求替换为RandomCrop等裁剪操作
transform = transforms.Compose([
    transforms.Resize((512, 512)),
    # 可补充归一化、增强等其他操作
])
dataset = AidaDataset(image_dir=IMG_DIR, mask_dir=MASK_DIR, transform=transform, target_transform=transform)

该方案无需修改DataLoader和训练逻辑，是语义分割训练最常用的方案。

内容的提问来源于stack exchange，提问作者Jonathan Woollett-light