3. 数据加载与预处理问题：Dataset与DataLoader实现细节、数据增强导致的形状不匹配、内存泄漏排查

数据加载这块，说实话，是联调时最容易出幺蛾子的环节。你想想看，模型结构调了半天，loss曲线也正常，结果一跑起来直接OOM，或者训练到一半形状对不上崩了——这种体验我太熟了。今天咱们就把这几个坑一个一个踩明白。

3.1 Dataset与DataLoader的实现细节

先聊聊Dataset。很多人觉得这玩意儿就是个“读图片+返回标签”的壳子，其实不然。我见过太多人把数据预处理全塞在__getitem__里，结果训练慢得像蜗牛。为什么？因为DataLoader是多进程预取的，但如果你在__getitem__里做大量计算，每个worker都在重复造轮子。

举个例子，我习惯这样写：

class MyDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.img_paths = sorted(glob.glob(root_dir + '/*.png'))
        self.labels = [self._parse_label(p) for p in self.img_paths]
        self.transform = transform
        # 提前加载一个缓存字典，避免重复读取
        self.cache = {}

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

    def __getitem__(self, idx):
        if idx in self.cache:
            img = self.cache[idx]
        else:
            img = cv2.imread(self.img_paths[idx])
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            self.cache[idx] = img  # 小数据集可以缓存，大数据集慎用

        label = self.labels[idx]
        if self.transform:
            img = self.transform(img)
        return img, label

这里有个小细节：self.cache只适合小数据集。我曾在项目里缓存了10万张图，结果内存直接爆炸。嗯，这个教训挺深刻的。

3.2 数据增强导致的形状不匹配

这是联调时最头疼的问题之一。你辛辛苦苦写好了数据增强流水线，结果一跑就报错：RuntimeError: The size of tensor a must match the size of tensor b。为什么会这样？

说白了，就是数据增强操作改变了张量的维度，但后续的collate函数没跟上。我遇到过最典型的情况：

• 随机裁剪后，同一batch里的图片尺寸不一样了
• 旋转操作改变了通道顺序，从CHW变成了HWC
• 某些增强操作（如MixUp）需要特殊处理标签

正确的做法是：在数据增强流水线的最后，加一个Resize或Pad操作，确保输出尺寸一致。比如：

from torchvision import transforms

train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

注意看，RandomResizedCrop(224)这里指定了输出尺寸为224x224。如果你不写这个参数，它默认是随机尺寸，那就等着报错吧。

3.3 内存泄漏排查

内存泄漏这个问题，说大不大，说小不小。训练到一半突然OOM，或者显存占用只增不减——我敢说每个做感知的人都遇到过。排查起来其实有套路。

先看几个常见原因：

现象	可能原因	排查方法
显存持续增长	计算图未释放、梯度累积	检查loss.backward()和optimizer.step()是否成对出现
CPU内存缓慢上升	DataLoader的num_workers泄漏	减少worker数，或使用torch.multiprocessing的共享内存
训练中途OOM	数据增强时创建了过多临时变量	用gc.collect()手动回收，或检查__getitem__里的循环

我个人习惯用torch.cuda.memory_summary()来查看显存分配情况。这个函数会打印出每个张量的地址和大小，非常直观。另外，nvidia-smi只能看个大概，真要定位问题还得靠代码里的torch.cuda.empty_cache()。

我曾经遇到过一个特别隐蔽的泄漏：在__getitem__里用了cv2.imread，但没释放OpenCV的Mat对象。OpenCV的Mat在Python里是自动管理的，但如果你在循环里创建了大量临时变量，GC来不及回收，内存就慢慢涨上去了。解决方案很简单：用with语句或者显式调用del。

3.4 知识体系总览

下面这张图是我自己总结的数据加载与预处理核心逻辑，你可以对照着看：

这张图把数据加载的整个链路画出来了。你看，从原始数据到模型训练，中间经过Dataset和DataLoader两个关键节点。每个节点都可能出问题，但归根结底就两件事：数据流的一致性（形状、类型、顺序）和资源管理（内存、显存）。

嗯，今天就先聊到这儿。数据加载这块其实还有很多细节，比如多进程的共享内存、pin_memory的用法、自定义collate_fn的写法等等。不过先把这三个最核心的问题搞定，联调时至少能少踩一半的坑。