PyG 加载 QM9 数据集从下载到跑通训练【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric第一次用 PyTorch GeometricPyG一个基于 PyTorch 的图神经网络库做分子属性预测时QM9 几乎是绕不开的入门数据集——它装着 13 万多个有机小分子和 19 个量子化学回归目标。可实际操作里下载依赖、目标字段切列、量纲归一化这三步经常卡住人。这篇文章的思路是先用最短的代码把数据跑起来再回过头解释每一步在做什么遇到报错时知道往哪儿查。十分钟把 QM9 加载进内存先跑起来再说。下面这段代码在任意能装torch-geometric的目录里直接执行首次运行会自动把原始文件下载到data/QM9/raw/再处理成data/QM9/processed/data_v3.ptfrom torch_geometric.datasets import QM9 dataset QM9(rootdata/QM9) print(len(dataset)) # 130831 个分子 print(dataset.num_features) # 11 维原子特征 print(dataset.num_targets) # 19 个回归目标 print(dataset[0].z.shape) # 第一个分子的原子序数第二次运行同一段代码时PyG 检测到processed/目录里已经有data_v3.pt就跳过下载与处理直接读盘通常几秒内返回。这个缓存到磁盘的行为来自QM9的父类InMemoryDataset具体逻辑在 torch_geometric/datasets/qm9.py 里可以看到。如果你只是想验证环境能不能跑通到这里就够了。下面我们把数据里到底装了什么拆开看。QM9 里每个分子长什么样dataset[i]返回的是一个Data对象把它想象成一行分子体检报告z每个原子的原子序数整数张量pos每个原子的三维坐标笛卡尔坐标QM9 只保留了能量最低的构象x11 维原子特征由 5 类原子 one-hotH/C/N/O/F加上原子序数、芳香性、sp/sp²/sp³ 杂化、连氢数拼起来y19 个回归目标从偶极矩、极化率一直到各种能量edge_index/edge_attr化学键连接关系和键型 one-hot。19 个目标对应的物理量和单位源码里用一张表格列得很清楚直接打开 torch_geometric/datasets/qm9.py 翻到QM9的 docstring 就能看到完整清单。做能量类目标U_0、U、H、G以及对应的 atomization 版本时通常要先减去原子参考能量再做归一化QM9.atomref(target)就是干这个的——它按原子序数查一张内置表把每个原子的孤立原子能量扣掉让网络只需学成键带来的能量变化收敛会快很多。下载与处理RDKit 决定走哪条路这里有个容易踩的坑PyG 的QM9内部会尝试import rdkit成功与否决定了它下载哪份原始文件、走哪条处理分支。有 RDKit下载qm9.zipSDF 结构 目标 CSV 一个 skip 列表然后自己解析 SDF 生成z / pos / x / edge_index / smiles。好处是能拿到 SMILES 字符串后续做 SMILES 相关模型时方便。无 RDKit改下qm9_v3.zip里面是官方预处理好的qm9_v3.pt。好处是不用装化学库坏处是数据里没有smiles字段其他字段与手动处理版一致源码process()里也做了对应分支。所以如果你后面要写涉及 SMILES 的代码建议先把 RDKit 装上conda install -c conda-forge rdkit # 或者pip install rdkit-pypi装完再触发一次处理QM9(root..., force_reloadTrue)会强制重跑download()和process()。如果你只想快速看效果、暂时不想装 RDKitPyG 会自动切到预处理好版本功能基本不受影响只是data.smiles是None。第一次跑的时候如果下载中断或者目录被手动删了一半可能会看到类似这样的报错FileNotFoundError: [Errno 2] No such file or directory: data/QM9/processed/data_v3.pt通常是raw/下的原始文件缺了一份比如uncharacterized.txt没下载完PyG 判断 raw 文件不齐就跳过了process()。处理方法是把data/QM9/整个目录清掉再重新执行QM9(rootdata/QM9)让它从头走一遍下载与处理。训练时最容易卡住的三处数据集加载本身只是第一步真正训练时几个字段格式问题反复出现。第一处y有 19 列但你的模型只预测 1 个。直接拿data.y喂给SchNet之类的回归头形状对不上。最简单的做法是在数据集构造时就用transform把目标列切出来参考 examples/qm9_nn_conv.py 里的MyTransformimport copy class KeepOneTarget: def __init__(self, target: int): self.t target def __call__(self, data): data copy.copy(data) data.y data.y[:, self.t:self.t1] return data把它塞进QM9(root..., transformKeepOneTarget(0))之后data.y就只有一列了。第二处不同目标量级差很多。偶极矩、极化率、HOMO 能量、自由能各自量纲和数值范围都不一样直接 MSE 训练时大数值目标会主导梯度。惯例是先做标准化减均值除标准差训练完评估 MAE 时再乘回去。官方示例也是这么做的dataset.data.y (dataset.data.y - mean) / std之后std要单独存下来给评估用。第三处DimeNet 系列模型的预训练示例里有一行列重排。如果你在 examples/qm9_pretrained_dimenet.py 里看到dataset.data.y dataset.data.y[:, idx]且idx [0,1,2,3,4,5,6,12,13,14,15,11]别困惑——那是因为 DimeNet 论文里把U_0/U/H/G列 7/8/9/10换成了对应的 atomization 版本列 12/13/14/15δe列 4又用e_LUMO - e_HOMO代替。跑 SchNet 或 NNConv 时不需要这一步只有加载 DimeNet 预训练权重时才要对齐列顺序。一个端到端的分子属性预测工作流把上面几步串起来就是一个最小可跑的 QM9 训练脚本。为了短这里用SchNet一个基于连续滤波器的分子 GNN输入是原子序数 坐标输出标量属性完整代码可以对照 examples/qm9_nn_conv.py 扩展。准备数据切目标列并做归一化然后按官方习惯做随机划分最后 1 万做测试前 1 万做验证其余训练import torch from torch_geometric.datasets import QM9 from torch_geometric.loader import DataLoader dataset QM9(rootdata/QM9, transformKeepOneTarget(0)).shuffle() y dataset.data.y mean, std y.mean(dim0, keepdimTrue), y.std(dim0, keepdimTrue) dataset.data.y (y - mean) / std train_ds dataset[10000:] val_ds dataset[:10000] train_loader DataLoader(train_ds, batch_size64, shuffleTrue) val_loader DataLoader(val_ds, batch_size64)模型和优化器SchNet只吃z / pos / batch三个字段坐标是分子图的关键输入别忘了data.pos必须在数据集里QM9 默认就带import torch.nn.functional as F from torch_geometric.nn import SchNet model SchNet(hidden_channels128, num_filters128, num_interactions6, num_gaussians50) opt torch.optim.Adam(model.parameters(), lr1e-3)训练循环评估时用std把归一化后的 MAE 还原回原始量纲for epoch in range(1, 21): model.train(); total 0.0 for batch in train_loader: opt.zero_grad() pred model(batch.z, batch.pos, batch.batch) loss F.mse_loss(pred.view(-1), batch.y.view(-1)) loss.backward() opt.step() total loss.item() * batch.num_graphs print(fepoch {epoch:02d} train loss {total/len(train_ds):.4f})跑几个 epoch 后 loss 应该明显下降。如果训练速度不够快可以把num_workers加到 2~4如果显存紧张先把hidden_channels/num_filters从 128 砍到 64。源码与示例入口数据集实现torch_geometric/datasets/qm9.pydownload()/process()/atomref()都在这个文件里分子 GNN 最小示例NNConv GRU Set2Setexamples/qm9_nn_conv.py用 SchNet 预训练权重跑 12 个目标的 MAE 评测examples/qm9_pretrained_schnet.py用 DimeNet / DimeNet 预训练权重评测含列重排示例examples/qm9_pretrained_dimenet.py更大规模的分子图数据PCQM4M分布式训练参考examples/multi_gpu/pcqm4m_ogb.py同类分子数据集 ZINC15 万分子1 个二分类目标torch_geometric/datasets/zinc.pyQM9 的关键就三件事先跑起来拿到Data对象再切出你要的那个目标列并归一化最后按分子图模型的输入习惯z / pos / batch或x / edge_index / batch把字段喂对。把这三步走顺换到 ZINC 或 PCQM4M 上只是换个数据集类名的事。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考