如何在不统一图像尺寸的情况下用PyTorch训练DeepLab模型？

针对你不想统一图像尺寸、又要避免单样本训练的需求，以下是几种PyTorch框架下的实用方法：

1. 让模型支持可变输入尺寸 #

DeepLab v3/v3+本身基于空洞卷积和编码器-解码器结构，天然支持可变尺寸输入，但要确保模型没有固定维度的约束：

• 移除或替换模型中所有固定尺寸的层（比如早期版本中可能存在的全连接层），改用AdaptiveAvgPool2d这类自适应池化层。
• 确认空洞卷积的设置不会限制输入尺寸（只要输入尺寸能被骨干网络的下采样倍数整除即可，比如ResNet50下采样32倍，输入尺寸最好是32的倍数，避免特征图尺寸出现小数）。

2. 自定义DataLoader的批次处理逻辑 #

利用PyTorch DataLoader的collate_fn参数，实现不强制统一尺寸的批次组装，再通过梯度累积模拟批量训练：

自定义collate_fn #

def variable_size_collate(batch):
    # 直接返回原始尺寸的图像和掩码列表，不做resize
    images = [sample[0] for sample in batch]
    masks = [sample[1] for sample in batch]
    return images, masks

训练循环中实现梯度累积 #

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = DeepLabV3Plus(num_classes=your_num_classes).to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# 初始化DataLoader，传入自定义collate_fn
dataloader = torch.utils.data.DataLoader(
    your_dataset,
    batch_size=8,  # 这里的batchsize是单次读取的样本数
    shuffle=True,
    collate_fn=variable_size_collate
)

accumulation_steps = 8  # 每累积8个样本的梯度再更新一次
optimizer.zero_grad()
step = 0

for images, masks in dataloader:
    for img, mask in zip(images, masks):
        img = img.unsqueeze(0).to(device)
        mask = mask.unsqueeze(0).to(device)
        
        output = model(img)
        loss = criterion(output, mask)
        loss = loss / accumulation_steps  # 损失归一化
        loss.backward()
        
        step += 1
        if step % accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()
            print(f"Step {step}, Loss: {loss.item() * accumulation_steps}")

这种方式本质是把多个样本的梯度累积后再更新优化器，效果和真实batchsize=accumulation_steps的批量训练一致，同时保留了图像的原始尺寸。

3. 多尺度分组批次训练 #

如果想利用GPU的并行计算效率，可以在数据集预处理时，把尺寸相近的图像归为一组，同一批次内使用相同尺寸，不同批次保留不同尺寸：

• 先统计所有图像的尺寸，按宽高范围分成若干组（比如小图组、中图组、大图组）。
• 自定义Sampler，让DataLoader每次从同一组中采样样本，这样同批次尺寸一致，不同批次尺寸不同，既保证批量并行，又不用统一所有图像尺寸。

关键注意事项 #

• 确保损失函数适配可变尺寸：PyTorch的CrossEntropyLoss、DiceLoss等常用分割损失，只要输出特征图和掩码的空间维度匹配，就能正常计算。
• 避免固定尺寸的数据增强：训练时只做随机翻转、颜色扰动、高斯模糊等不改变图像尺寸的增强操作，不要强制resize。

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