企业官网建设流程全解析

InMemoryDataset 深度指南：PyTorch Geometric 图数据加载与内存优化

【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric

在 PyTorch Geometric 里搭第一个自定义图数据集时，大多数人的第一反应是直接遍历一张一张图地存；可当样本量上去之后，内存被悄悄吃掉，训练还没开始就先爆掉。而 PyG 的InMemoryDataset走的是一条完全不同的路——它不把样本一张张存起来，而是先"合并"再存储。搞懂这一招，图数据加载和内存优化就不再是玄学，选型也有了依据。

一本大账本：合并存储为什么更省内存

先打个比方：把 10 万本各装几页的小账本，合并成一本厚账本，再附一张页码索引表，记下"第 N 本记在哪个区间"。图数据集的"合并存储"就是这个思路——所有样本的特征张量首尾拼接成一个大张量，slices字典记录每个样本在各属性上的起止位置（本质是累积和）。

落到代码里，核心就两件事：

• collate()（见torch_geometric/data/collate.py）把一组Data拼成单个大对象：特征走torch.cat，slices存拼接后各段的偏移量。注意edge_index这类"节点编号"属性不能直接拼——第二个图的节点 0 必须接着第一个图的编号走，所以合并时要先加上偏移量；
• 取回样本时，separate()（见torch_geometric/data/separate.py）按slices[idx]到slices[idx+1]用narrow切出区间，再把偏移量减回去，还原出原始编号。

省内存的逻辑在于：内存里只有"一份"数据而不是 N 份对象壳，且特征连续存放在大张量里，访问顺序友好。这也是为什么 Cora、CiteSeer 这类中小规模引用数据集非常适合放在内存里跑。

数据是怎么流动的：从落盘到批处理

一次完整的存取链路可以串成三句话。

写：你在process()里拿到data_list后调用collate()，得到合并后的大对象和slices，再由InMemoryDataset.save()把(data, slices, 类型)序列化到processed/目录——此后原始文件就可以只留作备查。

读：实例化时load()把大对象整块读进内存。第一次调用get(idx)时，separate()按索引切出第 idx 个样本，同时顺手存进内部列表_data_list；之后再取同一个 idx，直接命中缓存返回副本，不再重复切分。这个缓存对"同一张图反复取"的访问模式（比如验证集轮次访问）收益最明显。

用：DataLoader每个 worker 对数据集做dataset[idx]触发上述get()路径，随后Batch.from_data_list()把一个 mini-batch 里的样本再合并成批对象。讽刺的是，这里用的合并机制和"合并存储"是同一套collate——slices 在批处理环节还会被复用成batch/ptr向量。

不同数据规模下 Dataset 怎么选

InMemoryDataset和OnDiskDataset（torch_geometric/data/on_disk_dataset.py，底层支持sqlite/rocksdb两种后端）的定位差异，可以这样对比：

经验上：Cora 级别的引用网络、QM9 级别的分子集放心用内存版；当特征拼起来的总量逼近可用内存、或你准备多机训练时，就该考虑切到磁盘格式了。

内存吃紧时如何切换到磁盘存储

最省事的路径是官方提供的转换方法，一段代码即可完成：

import torch_geometric disk_ds = dataset.to_on_disk_dataset(root='data/molecule_on_disk', backend='sqlite')

它会把内存集逐条（每 1000 条一批）写进数据库文件。两点提醒：当前实现只支持同构图（异构图会抛NotImplementedError）；如果数据规模大到"先全量进内存再转出"都不现实，就跳过转换，直接继承OnDiskDataset重写process()，边读原始数据边extend()落库。

加载太慢时，先确认优化方向没走偏：

• 保证pre_transform的产物已落盘到processed/——它只应执行一次，后续实例化直接load()；
• 对固定子集（如训练索引）做dataset.copy(train_idx)，数据集本体变小，缓存也更好命中；
• 内存够的话，dataset.to('cuda')可把整份特征一次性搬上卡，省掉逐样本搬运的开销。

分布式训练下 InMemoryDataset 的替代思路

InMemoryDataset的全部数据都挂在单个进程里，多卡、多节点场景下没有"每个 worker 读自己那份"的能力，这是它的设计边界而非 bug。两条替代路线：

1. 转磁盘 + 官方分布式流程：转成OnDiskDataset后，按examples/distributed/里的方案做数据分区，采样器子进程负责拉取邻居，训练进程只消费自己分区内的数据。整图先按连通块划分到各机器，再各自采样：

2. 干脆不整图加载：用NeighborLoader/DistNeighborLoader这类采样式加载器按需取局部子图，全量图可以只存在于磁盘，内存压力从"样本总量"降到"单步邻域大小"。

完整实战：自定义 InMemoryDataset 模板

这段代码展示一个可运行的最小模板：子类只需声明两个文件名属性，并在download()/process()里各放一段逻辑，基类会自动处理目录结构、缓存判断与落盘。

import os.path as osp import torch from torch_geometric.data import Data, InMemoryDataset class MoleculeSet(InMemoryDataset): def __init__(self, root, transform=None, pre_transform=None): super().__init__(root, transform, pre_transform) # 基类发现 processed 文件缺失时，会自动走 download() -> process() merged, self.slices = torch.load(self.processed_paths[0]) self.data = merged @property def raw_file_names(self): return ['molecules_raw.pt'] @property def processed_file_names(self): return ['molecule_merged.pt'] def download(self): # 这里放原始数据的获取逻辑；若数据已就位，留空即可 torch.save(self._load_raw_source(), self.raw_paths[0]) def process(self): raw = torch.load(self.raw_paths[0]) # 逐样本清洗、补全，再交给 pre_transform（若有） graphs = [] for mol in raw: g = self._prepare(mol) if self.pre_filter is not None and not self.pre_filter(g): continue if self.pre_transform is not None: g = self.pre_transform(g) graphs.append(g) # 合并为单一大对象 + slices，一并落盘 merged, slices = self.collate(graphs) torch.save((merged, slices), self.processed_paths[0])

各方法的分工：两个*_file_names属性只回答"哪些文件构成数据"，基类据此判断是否需要重跑；download()负责把原始数据放进raw/；process()是真正的流水线——清洗、过滤、变换、合并、落盘。之后无论重启多少次进程，__init__里的那句torch.load都会直接命中processed/，合并开销只发生一次。

接入训练循环则交给DataLoader，它会按索引取样本并自动完成批合并：

from torch_geometric.loader import DataLoader loader = DataLoader(MoleculeSet('data/molecule/'), batch_size=32, shuffle=True) for batch in loader: out = model(batch.x, batch.edge_index, batch.batch)

三条可以直接执行的收尾建议

• 动手前算一笔账：样本特征总量（节点数 × 维度 × dtype 字节数）是否小于可用内存的 2/3，超了就直上磁盘格式，别和内存硬扛。
• 训练/验证/测试拆集用dataset.copy(idx)生成子集，而不是每次循环里都过滤原始集；子集更小，缓存命中率更高。
• 需要多机或大模型采样时，尽早把流程迁到OnDiskDataset或采样式加载器；InMemoryDataset留给单机、中规模、读多写少的场景。

机制层面的细节（collate/separate对异构图、稀疏张量的处理）与建集完整教程，建议对照仓库内docs/source/tutorial/的 create_dataset 一节和torch_geometric/data/目录下的源码注释精读一遍，配合本文的类比会容易得多。

【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric