
PyTorch Geometric 加载 QM9 数据集从环境自检到跑通一次分子属性预测【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric本文带你用 PyTorch Geometric 完成 QM9 数据集加载的完整链路先做环境与依赖自检再一次性把数据落盘然后核对数据、接入 SchNet 跑通一个训练 epoch最后给出批处理提速与排错清单。QM9 是 GNN 分子建模的标准测试集约 13 万个含 C、H、N、O、F 的有机小分子每个分子带 3D 原子坐标与 19 种量子化学性质偶极矩、HOMO/LUMO 能级、内能、焓等平均 18 个原子、37 条键。它也是 PyG 中InMemoryDataset内存数据集基类首次处理后把全部图存成一个.pt缓存之后再打开直接载入内存的典型使用者官方实现在 torch_geometric/datasets/qm9.py。QM9 环境自检RDKit 依赖与降级加载开始之前先确认两件事避免「PyG QM9 加载失败」类问题在训练中途才暴露。装好 PyG 后先在终端验证导入执行python -c import torch_geometric再单独执行python -c import rdkit。RDKit 是化学信息学工具包这里负责解析 QM9 的原始 SDF 文件一种存放分子结构与坐标的文本格式把分子转成原子、键、坐标等图字段。QM9 的处理流程对 RDKit 有依赖缺失时会静默降级download()尝试import rdkit失败时改从processed_url下载官方预处理包qm9_v3.pt跳过 SDF 解析process()检测到未安装时在 stderr 打印一条提示后直接载入预处理数据功能影响拿不到 SMILES 字符串与name等分子结构字段其余 19 列目标值不受影响正常训练可以继续。安装建议优先conda install -c conda-forge rdkit其次pip install rdkit-pypi偶有依赖冲突按需处理。 注意一个坑如果你曾在无 RDKit 状态下加载过 QM9processed/data_v3.pt缓存里存的就是降级版本。装上 RDKit 后不会自动重新处理需要删掉该数据集目录或构造时传force_reloadTrue触发重处理。QM9 一次加载路径配置与 FileNotFoundError 排查QM9 构造时只需给一个root目录下载与处理全自动完成from torch_geometric.datasets import QM9 dataset QM9(root./data/QM9) # 首次运行自动下载并处理 print(len(dataset), dataset.data.y.shape)目录约定raw/存放原始文件装了 RDKit 是gdb9.sdf等没装是qm9_v3.ptprocessed/生成data_v3.pt缓存。构造完成后数据已整体载入内存再次实例化时直接读缓存重开成本从分钟级降到秒级需要重新处理时传force_reloadTrue即可。路径配置建议与排错提示不要照抄示例里基于__file__的相对拼接路径Jupyter 中__file__不可用osp.dirname(osp.realpath(__file__))会拿到非预期位置。改用显式的绝对路径或相对当前工作目录的./data/QM9。报FileNotFoundError: .../raw/gdb9.sdf多为 root 写权限不足、目录被提前删了一半或断网导致下载中断。检查目录可写、补全下载或手动把文件放进raw/。报 RDKit 相关ImportError不是崩溃而是走了上述降级路径装上 RDKit 再按提示重建缓存即可。数据核对19 列目标与 y 属性索引错位排查加载后先核对规模与形状再决定用哪一列训练import torch print(len(dataset)) # 130831 print(dataset.data.y.shape) # torch.Size([130831, 19]) # DimeNet 预训练权重对应原子化能量列需要先重排 y idx torch.tensor([0, 1, 2, 3, 4, 5, 6, 12, 13, 14, 15, 11]) dataset.data.y dataset.data.y[:, idx]19 列依次为偶极矩 μ、极化率 α、HOMO 能量、LUMO 能量、能隙 Δε、电子空间展布、ZPVE、内能 U₀/U、焓 H、自由能 G、热容 cv以及四组原子化能量U₀ᴬᵗᵒᵐ、Uᴬᵗᵒᵐ、Hᴬᵗᵒᵐ、Gᴬᵗᵒᵐ和三个转动常数。处理时 PyG 已把 Hartree、kcal/mol 等原始单位统一换算成 eV你拿到的就是可训练的量纲。⚠️ 索引错位问题DimeNet 官方预训练权重是按原子化能量列训练的所以 examples/qm9_pretrained_dimenet.py 先把 y 的第 7–10 列U₀、U、H、G替换成第 12–15 列对应原子化能量再把 cv 挪到第 11 位。若不重排就传target索引预测值会跟错列评估 MAE 明显异常甚至越界报错。另一条捷径是工厂方法SchNet.from_qm9_pretrained(path, dataset, target)会自行下载权重、切分 train/val/test参考 examples/qm9_pretrained_schnet.py。注意调用它时传原始未重排的 dataset索引逻辑由方法内部处理先重排再传target反而会错位。接入训练QM9 分子属性预测最小主线下面以偶极矩第 0 列为目标的 SchNet 训练为主线完整版见 examples/qm9_nn_conv.pyimport torch.nn.functional as F import torch_geometric.transforms as T from torch_geometric.datasets import QM9 from torch_geometric.loader import DataLoader from torch_geometric.nn import SchNet dataset QM9(./data/QM9, transformT.Distance(normFalse)).shuffle() data dataset[0] data.y data.y[:, 0] # 只保留目标列其余同理 train, val, test dataset[:100000], dataset[100000:110000], dataset[110000:] loader DataLoader(train, batch_size64, shuffleTrue) model SchNet(hidden_channels128, num_filters128, num_interactions6, num_gaussians50).to(cuda) opt torch.optim.Adam(model.parameters(), lr1e-3) for epoch in range(1, 3): for data in loader: data data.to(cuda) opt.zero_grad() out model(data.z, data.pos, data.batch) loss F.mse_loss(out.view(-1), data.y) loss.backward(); opt.step()要点说明T.Distance在线计算原子对距离特征SchNet 的 forward 只吃z原子序数、pos坐标与batch批内分子归属dataset.shuffle()后按位置切片做划分是官方示例的标准做法追求精度可对目标做均值/标准差归一化评估时再乘回去QM9 还提供dataset.atomref(target)原子参考能量用于把预测校正到原子化能量基线进阶再试。DataLoader 批处理提速与 pre_transform 预处理缓存动态批处理分子大小有差异时固定batch_size的批次内节点数会波动。改用DynamicBatchSampler(dataset, max_num_nodes1000)并传给DataLoader(dataset, batch_samplersampler)按节点数上限组批单批显存更平稳训练吞吐通常更高。pre_transform 落盘缓存transform每次访问数据都会重新执行pre_transform只在首次处理时执行一次结果直接写进processed/data_v3.pt。凡是不依赖随机性、每次都要做的特征加工如只保留某一目标列都应放进pre_transform后续加载零成本。多进程加载可设num_workers413 万分子全量加载约占数 GB 内存机器紧张时先切片到子集验证流程。排错清单与延伸练习把前文分散的排错提示汇总成自检清单报FileNotFoundError且路径含raw/gdb9.sdf确认 RDKit 状态与目录权限必要时手动补文件stderr 出现 Using a pre-processed version of the dataset属正常降级但装好 RDKit 后需重建缓存才能拿到 SMILES评估时 target 索引越界或 MAE 异常检查是否做了 y 列重排、target是否落在 [0, 11]重开 notebook 数据「不见了」确认processed/data_v3.pt是否存在存在则加载应只需秒级内存不足用切片子集先跑通再上全量。延伸练习把 SchNet 换成 DimeNetPlusPlus走一遍from_qm9_pretrained y 列重排的完整流程用dataset.atomref(12)对原子化能量目标做参考能量校正比较校正前后 MAE把数据接入LightningDataModule见 torch_geometric/data/lightning/或直接换用 PCQM4M 数据集体会预训练迁移效果。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考