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 Geometric（PyG）里，InMemoryDataset的"合并存储 + slices 切片"就是图数据集省内存的标配做法。这篇 10 分钟讲清选型标准、运行机制和迁移路线，代码只留能跑通的最小份。

⚖️ 先选型：InMemoryDataset 到底适不适合你

结论先行：整份数据集处理后的单文件小于机器一半内存，选 InMemory；否则直接上 OnDiskDataset，别纠结。

一个可操作的估算式：图数量 × 平均节点数 × 特征维度 × dtype 字节数（float32 按 4B 算）。比如 1 万张图、均 500 节点、64 维 → 约 128MB，随便放；但 10 万张图、均 5000 节点、1024 维 → 20GB+，就该走磁盘了。

还有个特殊情况：你的数据本身就是一张图（如 ogbn-products 的引用网络）。此时len()直接返回 1，get()走copy.copy(self._data)快路径，in_memory_dataset.py 第 103～104 行的设计就是为它准备的，用它没毛病。

🔎 机制速览：为什么"合并存储 + slices"能省内存

一句话：它把 N 个Data对象打包成一个Data大对象，外加一份"分装清单"slices——相当于把 N 个包裹打成一个快递箱，箱里贴好"第几件占第几段"的清单，取货时照单裁切。

省在哪？每个独立的Data都是一份 Python 对象头 + 若干属性 dict + 每个张量独立的 storage 元数据与对齐填充；pickle 落盘时这份开销逐图重复。合并后，同构属性变成一块连续大张量，清单本身只存N+1个下标。

存取链路对应源码三处：

1. 存：process()里调self.collate(data_list)（底层是 collate.py，increment=False, add_batch=False），产出(data, slices)存盘；
2. 取：get(idx)调 separate.py 的separate()，对每个属性做narrow(attr, slices[idx], slices[idx+1])，把大张量裁回小图；
3. 缓存：首次取到的样本存进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'][i+1])

拿 1000 张均 500 节点、64 维 float32 的图举例：x合并成[320000×64]一块连续 tensor，约 8MB；slices三个属性各占 1001 个 int64，总共不到 24KB。而 N 份独立对象 pickle 后的文件体积通常比合并后大出 3～4 倍——这就是下面账单里 70%+ 节省的来源。

🔧 跟着敲：4 步跑通一个最小的 InMemoryDataset

结论：只需要__init__里加载、两个文件名属性、process里collate落盘，四步走完。下面这个例子自己"生成"原始数据，本地即可跑：

import torch from torch_geometric.data import InMemoryDataset, Data class MyShapes(InMemoryDataset): # ① __init__：构造时把合并数据读回内存 def __init__(self, root, transform=None, pre_transform=None): 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(x=x, edge_index=edge_index, y=y)) 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(root='data/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_size=64, shuffle=True) # 自动把小图拼成 Batch

两个细节别踩坑：

• pre_transform只在process()时跑一次并固化进磁盘；transform是每次dataset[i]现场变换。想换处理逻辑，删掉processed/目录（或构造时传force_reload=True）。
• 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 RAM，PyG 2.x；数字作量级参考，你的机器按同一比例缩放即可）：

Reddit 这行的教训：合并本身要一份全量 RAM，超过内存上限时不是"慢"，是"起不来"。所以选型估算要放在写process()之前做。

🛠️ OOM 速查与迁移路线

故障速查表：现象 - 原因 - 解法

三步迁移：InMemory → OnDisk → 分布式

数据长大后不用重写数据集，走一条平滑路线：

1. 转存：一行on_disk = dataset.to_on_disk_dataset(root='data/on_disk', backend='sqlite')。它会自动从get(0)推断 schema（dtype + 拼接维度），每 1000 张图批量extend()写入，避免逐图写库。
2. 换加载器：训练侧把DataLoader换成DistLoader/ 采样器，图数据按分区从磁盘数据库流入训练进程，各节点只持有一部分：

3. 看整机接线：多机场景下图先按机器分区存放（跨机边在采样时远程取回）：

采样阶段的本地/远程节点划分长这样：

延伸资源

• 官方教程（创建数据集全流程）：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