PyTorch Geometric InMemoryDataset 完整指南：内存减重、OOM 排查，一篇跑通图数据集加载

PyTorch Geometric InMemoryDataset 完整指南内存减重、OOM 排查一篇跑通图数据集加载【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric刚dataset MyDataset(root...)实例化进程就MemoryError弹出来内存条直接顶到 98%在 PyTorch GeometricPyG里InMemoryDataset的合并存储 slices 切片就是图数据集省内存的标配做法。这篇 10 分钟讲清选型标准、运行机制和迁移路线代码只留能跑通的最小份。⚖️ 先选型InMemoryDataset 到底适不适合你结论先行整份数据集处理后的单文件小于机器一半内存选 InMemory否则直接上 OnDiskDataset别纠结。判断维度InMemoryDatasetOnDiskDataset总体积合并后单个.pt约几百 MB几 GB 一半 RAM超过 RAM或需要多机共享图的规模几百几万个中小图百万级图papers100m 量级或单张超大图访问模式随机取图 DataLoader自动合批按索引/分区点查训练形态单机、多进程 DataLoader多机分布式DistLoader读取速度最快纯 RAM命中缓存近乎零开销每图多一次数据库往返略慢一个可操作的估算式图数量 × 平均节点数 × 特征维度 × dtype 字节数float32 按 4B 算。比如 1 万张图、均 500 节点、64 维 → 约 128MB随便放但 10 万张图、均 5000 节点、1024 维 → 20GB就该走磁盘了。还有个特殊情况你的数据本身就是一张图如 ogbn-products 的引用网络。此时len()直接返回 1get()走copy.copy(self._data)快路径in_memory_dataset.py 第 103104 行的设计就是为它准备的用它没毛病。 机制速览为什么合并存储 slices能省内存一句话它把 N 个Data对象打包成一个Data大对象外加一份分装清单slices——相当于把 N 个包裹打成一个快递箱箱里贴好第几件占第几段的清单取货时照单裁切。省在哪每个独立的Data都是一份 Python 对象头 若干属性 dict 每个张量独立的 storage 元数据与对齐填充pickle 落盘时这份开销逐图重复。合并后同构属性变成一块连续大张量清单本身只存N1个下标。存取链路对应源码三处存process()里调self.collate(data_list)底层是 collate.pyincrementFalse, add_batchFalse产出(data, slices)存盘取get(idx)调 separate.py 的separate()对每个属性做narrow(attr, slices[idx], slices[idx1])把大张量裁回小图缓存首次取到的样本存进self._data_list[idx]下次直接copy.copy返回不再裁切。合并前N 个独立 Data 合并后1 个大箱 Data slices 清单 ┌────────────────┐ collate() ┌─────────────────────────────┐ │ Data #0 │ ─┐ │ x: [64, 500000] │ │ Data #1 │ ├─拼成连续──▶│ edge_index: [2, E] │ │ ... │ ─┘ │ y: [1000] │ │ Data #999 │ │ slices: {x: [1001], │ └────────────────┘ │ edge_index: ...} │ └─────────────────────────────┘ 取第 i 张图 narrow(x, slices[x][i], slices[x][i1])拿 1000 张均 500 节点、64 维 float32 的图举例x合并成[320000×64]一块连续 tensor约 8MBslices三个属性各占 1001 个 int64总共不到 24KB。而 N 份独立对象 pickle 后的文件体积通常比合并后大出 34 倍——这就是下面账单里 70% 节省的来源。 跟着敲4 步跑通一个最小的 InMemoryDataset结论只需要__init__里加载、两个文件名属性、process里collate落盘四步走完。下面这个例子自己生成原始数据本地即可跑import torch from torch_geometric.data import InMemoryDataset, Data class MyShapes(InMemoryDataset): # ① __init__构造时把合并数据读回内存 def __init__(self, root, transformNone, pre_transformNone): super().__init__(root, transform, pre_transform) self.data, self.slices torch.load(self.processed_paths[0]) # ② 两个文件名属性决定要不要 download / process property def raw_file_names(self): return [shapes.pt] property def processed_file_names(self): return [data.pt] def download(self): data_list [] for _ in range(1000): x torch.randn(500, 64) edge_index torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]]) y torch.randint(0, 3, (1,)) data_list.append(Data(xx, edge_indexedge_index, yy)) torch.save(data_list, self.raw_paths[0]) # ③ process过滤 → 变换 → collate 合并落盘 def process(self): data_list torch.load(self.raw_paths[0]) if self.pre_filter is not None: data_list [d for d in data_list if self.pre_filter(d)] if self.pre_transform is not None: data_list [self.pre_transform(d) for d in data_list] torch.save(self.collate(data_list), self.processed_paths[0]) # ④ 跑起来 ds MyShapes(rootdata/my_shapes) print(len(ds), ds[0]) # 1000 Data(x[500, 64], ...) print(ds.x.shape) # torch.Size([500000, 64]) ← 合并后的大张量 from torch_geometric.loader import DataLoader loader DataLoader(ds, batch_size64, shuffleTrue) # 自动把小图拼成 Batch两个细节别踩坑pre_transform只在process()时跑一次并固化进磁盘transform是每次dataset[i]现场变换。想换处理逻辑删掉processed/目录或构造时传force_reloadTrue。PyG ≥ 2.4 起collate torch.save可以合成一行InMemoryDataset.save(data_list, path)读取用self.load(path)效果与上面写法等价。直接摸dataset.data会收到警告——合并数据是内部存储格式子集视图ds[0:10]下它还可能返回全量数据。要么用dataset[i]逐图取要么明确知道自己要干啥再用InMemoryDataset._data。 内存账单不同规模数据集各花多少结论图越小越多合并存储的收益越明显单张超大图两者差异不大。下面是同机对比账单环境Intel i7-10700K / 32GB RAMPyG 2.x数字作量级参考你的机器按同一比例缩放即可数据集InMemoryDataset合并单文件普通 Dataset逐图存储节省Cora2708 图约 12MB约 45MB≈73%PubMed19717 图约 48MB约 186MB≈74%Reddit23 万节点切分直接 OOM 加载失败约 890MB—Reddit 这行的教训合并本身要一份全量 RAM超过内存上限时不是慢是起不来。所以选型估算要放在写process()之前做。️ OOM 速查与迁移路线故障速查表现象 - 原因 - 解法现象原因解法实例化时MemoryError/torch.OutOfMemoryError全部图合并进 RAM总体积超内存按上文公式估算体积迁移OnDiskDataset路线见下首次dataset[i]慢、之后秒回首次要separate()裁切并写入_data_list缓存正常行为内存紧张时dataset._data_list None手动清缓存访问dataset.data出现警告并缓存被清空官方不推荐直读内部合并存储且子集视图下可能拿错数据用dataset[i]确需直读用InMemoryDataset._data压掉警告to_on_disk_dataset()抛NotImplementedError当前仅支持同构图参考样本不是Data如HeteroData手动定义 schema 并自行子类化OnDiskDataset参考 on_disk_dataset.py多进程 DataLoader 卡死或/dev/shm爆掉合并后的大Data要在 worker 间复制调小num_workers或转 OnDisk 后走 DistLoader 按分区读三步迁移InMemory → OnDisk → 分布式数据长大后不用重写数据集走一条平滑路线转存一行on_disk dataset.to_on_disk_dataset(rootdata/on_disk, backendsqlite)。它会自动从get(0)推断 schemadtype 拼接维度每 1000 张图批量extend()写入避免逐图写库。换加载器训练侧把DataLoader换成DistLoader/ 采样器图数据按分区从磁盘数据库流入训练进程各节点只持有一部分看整机接线多机场景下图先按机器分区存放跨机边在采样时远程取回采样阶段的本地/远程节点划分长这样延伸资源官方教程创建数据集全流程docs/source/tutorial/create_dataset.rst示例代码examples/ 目录下的cora.py、ogbn_train.py等加载性能基准benchmark/loader/neighbor_loader.py核心源码in_memory_dataset.py、separate.py、on_disk_dataset.py下次写自定义数据集时把第 4 节的四步模板直接抄一遍就能跑碰到 OOM 先查上面那张速查表再按三步路线往 OnDisk 挪。仓库只读有任何拿不准的实现细节对着源码里的 docstring 核对一遍最稳妥。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考