BSD100数据集实战：如何高效处理超分辨率图像训练任务

1次阅读

没有评论

共计 2008 个字符，预计需要花费 6 分钟才能阅读完成。

刚开始用 BSD100 做超分训练时，我的显卡总是爆显存。这个数据集虽然只有 100 张高质图片，但实际训练时会遇到三个头疼问题：

内存黑洞：把全部 288×288 的原图加载到内存，8GB 显存的卡连第一个 epoch 都跑不完
IO 瓶颈：传统预处理把图片放大 2x/4x 后存临时文件，SSD 都被读写出火星子了
计算浪费：75% 的训练时间花在等数据上，GPU 利用率像过山车

最夸张的时候，训练速度只有 1.2 samples/sec，还不如我在 CPU 上跑 …

试过两种常见的处理方案：

传统整图加载
预处理阶段把所有图片放大到目标尺寸
训练时整个 batch 直接塞进显存
优点：实现简单
缺点：显存占用是原始图的 16 倍（4x 超分时）
动态分块加载
只加载当前 batch 需要的图像区域
配合多进程预读取
优点：显存占用恒定
缺点：需要处理复杂的边界对齐

最终选择基于 PyTorch Dataloader 的 动态裁剪方案，因为：

原生支持多进程数据加载
可以无缝对接混合精度训练
社区生态完善（后续加 LR 调度器方便）

关键技巧在于 __getitem__ 方法里现算裁剪坐标：

class BSD100Dataset(Dataset):
    def __getitem__(self, idx):
        # 随机生成裁剪坐标
        y = random.randint(0, orig_h - crop_size)
        x = random.randint(0, orig_w - crop_size)

        # 用 Pillow 的懒加载特性
        with Image.open(self.paths[idx]) as img:
            crop = img.crop((x, y, x+crop_size, y+crop_size))

        # 转为 Tensor 时才真正读入内存
        return ToTensor()(crop)

用 torch.cuda.amp 两行代码就能启用：

scaler = GradScaler()  # 防止梯度下溢

with autocast():
    output = model(inputs)
    loss = criterion(output, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

把 ColorJitter 移到 GPU 上执行，速度提升 3 倍：

aug = transforms.Compose([
    transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)
    ], p=0.8),
    transforms.RandomGrayscale(p=0.2)
]).cuda()  # 关键在这！

这个代码片段包含了所有优化技巧：

# 配置数据加载
train_loader = DataLoader(BSD100Dataset(..., crop_size=128),
    batch_size=16,
    num_workers=4,
    pin_memory=True  # 减少 CPU 到 GPU 传输时间
)

# 混合精度训练循环
for epoch in range(100):
    for inputs, targets in train_loader:
        inputs = inputs.to('cuda', non_blocking=True)
        targets = targets.to('cuda', non_blocking=True)

        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, targets)

        # 梯度累积技巧
        if (i+1) % 2 == 0:
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

测试环境：RTX 3080 (10GB) + Ryzen 7 5800X