图数据加载把内存撑爆了PyG InMemoryDataset 的 4 个高频问题一次讲清【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric用 PyTorch GeometricPyG做图数据加载时InMemoryDataset几乎人人都会碰到它把整个图数据集合并进一块 CPU 内存里训练取图又快又稳可一旦图多、节点大collate阶段就能直接把内存打爆。下面挑了 4 个最高频的问题从存储原理一路讲到分布式适配帮你把这条路走顺。为什么InMemoryDataset不存 N 份图而是拼成一整块普通Dataset会保留 N 个独立的Data对象每张图一个每个对象都拖着一份 Python 元数据开销不小。InMemoryDataset换了个思路调用collate把所有图拼成一个大的Data对象同时用一个slices字典记下每张图在各字段里的起止位置。存储阶段只做了这一件事# save() 的核心合并所有图 记录切片边界 data, slices cls.collate(data_list) # 拼成一整块 fs.torch_save((data.to_dict(), slices, data.__class__), path)collate内部对每个字段做拼接节点特征沿批次维竖向堆叠、边索引各自平移最终返回两个东西拼好的data和记录边界的slices。这样磁盘上只存一份大张量而不是 N 份对象元数据开销几乎被抹平。完整实现在 in_memory_dataset.py。 这就是PyG collate原理的核心拼接换空间切片换索引。对 Cora、PubMed 这类中小数据集省下的内存非常可观代价是整块数据必须一次装进内存图一大就会出事——正是下一个问题。取一张图时发生了什么collate和separate是镜像操作collate负责拼separate负责拆两者是镜像关系。你调用dataset[idx]时get(idx)会沿着slices把第idx张图从大张量里切出来# get(idx) 的核心逻辑 if self._data_list[idx] is not None: # 命中缓存直接返回副本 return copy.copy(self._data_list[idx]) data separate(clsself._data.__class__, batchself._data, idxidx, slice_dictself.slices, decrementFalse) self._data_list[idx] copy.copy(data) # 首次访问后缓存 return dataseparate做的事很机械读slices[idx]和slices[idx1]拿到起止位置用narrow沿对应维度切片得到单张图。注意decrementFalse——collate时已用incrementFalse拼过、边索引没做全局偏移所以拆出来不用再减回这一步比正常 batch 拆分还轻一点。⚠️ 这里藏着一个常被忽略的点_data_list是逐样本缓存。训练多轮后访问过的图会在内存里各留一份副本等于在拼好的一整块之外又叠了一层torch_geometric 数据加载慢时内存占用会悄悄翻倍。想释放可dataset._data_list None下次再按需separate。内存不够用了从InMemoryDataset迁到磁盘存储的三步操作当数据集大到装不进 CPU 内存典型的InMemoryDataset内存溢出官方给了一条现成路径to_on_disk_dataset()。它把每张图序列化进sqlite或rocksdb改成按需读盘# 三步建内存数据集 - 转磁盘数据集 - 指定 backend默认 sqlite dataset MyDataset(rootdata/) # 1. 先按内存数据集跑通 on_disk dataset.to_on_disk_dataset( # 2. 转成 OnDiskDataset rootdata/on_disk, backendsqlite, # 3. 显式传 root 防串数据 )迁移时有几个坑要提前知道只支持同构图内部拿第 0 张图检查是不是DataHeteroData直接抛NotImplementedError异构图得自己写OnDiskDataset子类。schema决定效率它会用第 0 张图推断字段 schema填了dtype/size就能跳过 Python pickle读写快不少。root必须显式传不传会落到root/on_disk多份训练切分时容易串数据。OnDiskDataset的完整实现在 on_disk_dataset.py它的processed_file_names就是一个{backend}.db。磁盘读取天然比内存慢但它把全部图常驻的硬限制拆成了单图按需加载是大规模数据集绕不开的一步。多卡分布式训练时 Dataset 该怎么处理单进程用InMemoryDataset很爽可一上多卡、多节点整块数据只有一份、塞在某个进程内存里这个前提就崩了——每个训练进程都得独立取到自己的样本。原生InMemoryDataset不支持分布式读取硬套会让进程间抢同一块内存。PyG 的分布式方案是分区 采样解耦先把数据集切成分区各存一份OnDiskDataset采样交给独立子进程再通过队列把SamplerOutput送进各节点的DistLoader最后喂给对应 Model。数据分区、采样、训练三者互不阻塞。优先走磁盘分布式基本等于InMemoryDataset转OnDiskDataset加分区别指望整块数据共享给所有进程。参考官方示例examples/distributed/下有 GraphLearn、Kuzu、PyG 原生三套可跑的分布式采样实现照着搭最稳。图数据集 分布式训练的关键不在 Dataset 本身而在谁来采样、怎么把邻居子图送到模型Dataset 只是提供分区入口。中小数据集放心用InMemoryDataset记牢collate/separate这对镜像操作和_data_list缓存一旦内存吃紧或要上多卡就走to_on_disk_dataset()加分区采样这条路。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
