简介这是一套面向分子能量预测任务的图神经网络完整实现定位于个人学习与算法复现适合化学信息学、计算化学及深度学习初学者参考。资源围绕有机小分子数据集构建了以原子为节点、化学键为边的图卷积网络涵盖数据预处理、模型构建、训练优化、结果验证等模块并借助PyTorch Geometric完成消息传递与特征聚合。目前已有83人学习压缩包共33个文件、约5.62MB其中8个Python脚本负责数据加载、训练与推理7个CSV提供分子能量样本pt文件存放模型权重mol文件描述分子结构另有png结果图、说明文档和备份文件目录结构清晰。通过学习可掌握分子图表示、能量回归建模、特征重要性分析与误差可视化方法。配套代码与数据可直接复现完整预测流程便于学习交流与后续扩展。1. 分子能量预测为什么值得用 GNN从一张图的表示到量子化学目标把分子能量预测做成一个 GNN 项目很多人的第一反应是“这会不会太学术了”。实际落地时你会发现QM9 这种分子性质数据集里常规机器学习方法——比如用 RDKit 算一堆描述符喂给随机森林——也能到几十 meV 的误差但模型的表现高度依赖特征工程换一批分子就失效。图神经网络GNN的价值在于把“分子是什么”这个问题的答案直接交给网络原子是节点、键是边消息传递机制天然贴合化学中的电子相互作用。你不再需要手工设计特征数据集里带什么字段模型就读什么。这套 Python 源码加数据集的组合适合两类人一类是想在材料/药物虚拟筛选里搭一个可用的能量预测基线另一类是准备把 GNN 从玩具数据集迁移到自研分子数据的算法工程师。接下来我按自己搭这套系统的顺序把环境、数据、模型和踩过的坑完整过一遍。2. 分子图构建的核心为什么原子是节点、化学键是边2.1 量子化学目标与图表示的对应关系分子能量预测预测的不是某个原子的属性而是整个体系的标量能量。这个目标决定了模型必须把局部交互聚合成全局表征。化学里能量可以粗略拆成键能、角度张力、非键相互作用而 GNN 的消息传递正好是一层一层把原子邻居的信息聚合到中心节点再通过读出函数变成全局向量。这种对应关系不是巧合而是图神经网络最早被用在分子性质预测上的根本原因。QM9 数据集里每个分子有三种可用表示SMILES 字符串、3D 坐标、以及原子/键的显式列表。实操中我推荐用原子序数加化学键类型作为主特征3D 坐标可以留作后续扩展。原子序数决定元素种类键类型决定相互作用强度这两者已经能支撑一个基线模型跑出合理结果。坐标信息虽然包含键长和键角但它把 GNN 从拓扑模型变成了几何模型训练成本和对噪声的敏感度都会上一个台阶后面单独讨论。2.2 最小分子图构建代码构建分子图最省事的路径是 RDKit 解析 SMILES再把原子和键映射到 PyTorch Geometric 的 Data 对象。下面这段代码是整套系统里最核心的预处理函数我建议你直接抄到项目里的graph_builder.py。import torch from rdkit import Chem from rdkit.Chem import AllChem from torch_geometric.data import Data ATOM_ORDER [1, 6, 7, 8, 9] # H, C, N, O, F BOND_TYPE { Chem.rdchem.BondType.SINGLE: 0, Chem.rdchem.BondType.DOUBLE: 1, Chem.rdchem.BondType.TRIPLE: 2, Chem.rdchem.BondType.AROMATIC: 3, } def smiles_to_graph(smiles: str) - Data: mol Chem.MolFromSmiles(smiles) if mol is None: raise ValueError(fInvalid SMILES: {smiles}) mol Chem.AddHs(mol) # 显式加上氢原子QM9能量包含氢的贡献 atoms [] for atom in mol.GetAtoms(): atomic_num atom.GetAtomicNum() if atomic_num in ATOM_ORDER: one_hot [1.0 if atomic_num a else 0.0 for a in ATOM_ORDER] else: # 超出常见五元素集合的原子归入其他 one_hot [0.0] * len(ATOM_ORDER) [1.0] atoms.append(one_hot) x torch.tensor(atoms, dtypetorch.float) src, dst, edge_attr [], [], [] for bond in mol.GetBonds(): i bond.GetBeginAtomIdx() j bond.GetEndAtomIdx() src.append(i) dst.append(j) src.append(j) dst.append(i) bond_type BOND_TYPE.get(bond.GetBondType(), 3) edge_attr.append([1.0 if bond_type t else 0.0 for t in range(4)]) edge_attr.append([1.0 if bond_type t else 0.0 for t in range(4)]) edge_index torch.tensor([src, dst], dtypetorch.long) edge_attr torch.tensor(edge_attr, dtypetorch.float) return Data(xx, edge_indexedge_index, edge_attredge_attr)这段代码的关键点有三个。第一Chem.AddHs必须调用不带氢的分子图会让能量误差明显变大具体数据见第 5 章。第二边是无向的但 PyTorch Geometric 的消息传递要求显式构造双向边所以每条键在edge_index里出现两次。第三edge_attr用的是键类型 one-hot维度固定为 4对应单键、双键、三键、芳香键RDKit 对芳香键有自己的枚举值必须映射否则 one-hot 维度会错位。如果你后续要加入键长只需要把距离拼到这个向量的第 5 位模型结构不用动。这里还做了一个原子序数到 one-hot 的映射ATOM_ORDER里只有 H、C、N、O、F。QM9 的分子恰好只含这五种元素但解析外部分子时可能遇到 S、P 等我统一归到“其他”类保证维度不炸。如果你的目标数据集包含卤素记得把ATOM_ORDER扩充并同步调整 one-hot 长度否则后面模型里in_channels就会对不上。3. 环境准备与数据集落盘跑通项目的第一道关卡3.1 Python 环境与核心依赖这套系统跑通的门槛不在模型而在依赖环境。我建议你直接用 Python 3.9 以上的虚拟环境隔离不要在全局环境里硬装因为 PyTorch Geometric 的版本组合很敏感。以下是我每次搭环境都会用的命令顺序python -m venv gnn_env source gnn_env/bin/activate pip install --upgrade pip pip install torch torchvision pip install torch_geometric pip install rdkit numpy pandas安装完成后马上验证 PyTorch Geometric 是否可用这一步能避开后续最神秘的报错python -c import torch_geometric; print(torch_geometric.__version__)如果 import 阶段就报错多半是 CUDA 版 PyTorch 与 PyTorch Geometric 的扩展没对上。常见做法是把 torch 卸载后用 CPU 版重新装或者按 PyTorch Geometric 官方给出的 wheel 索引补装torch-scatter和torch-sparse。我自己的经验是先试纯 pip 安装跑通 CPU 版再上 GPU一上来就装全套编译版容易在配置上花掉半天。3.2 数据集的下载、解析与完整性检查PyTorch Geometric 把 QM9 数据集封装成了现成的类这一步非常省心。数据包含约 13 万个分子每个分子有原子属性、3D 坐标、以及多项量子化学标签。用下面这段代码把数据下载到本地from torch_geometric.datasets import QM9 import torch ROOT ./data/qm9 dataset QM9(rootROOT) print(f分子数量: {len(dataset)}) print(f特征维度: {dataset.num_node_features}) print(f标签维度: {dataset.num_target_properties}) data dataset[0] print(data)我第一次跑这段代码时输出很意外数据集的num_node_features显示的是 11 或 15 之类的维度并不是我们自定义的 one-hot。原因是 PyTorch Geometric 内置的 QM9 类自带特征处理器它会自动生成节点特征不经过我们上一步写的smiles_to_graph。这本身没问题但你要清楚哪些字段是内置的、哪些是手工加的。内置数据对象的y是一个向量里面包含多项能量与热力学性质我们预测目标取第 0 列即可那就是 U0基态内能。数据集下载完成后建议做一个完整性检查import torch subset [dataset[i] for i in range(1000)] num_nodes torch.tensor([d.num_nodes for d in subset]) num_edges torch.tensor([d.edge_index.size(1) for d in subset]) print(f节点数: min{num_nodes.min()}, max{num_nodes.max()}, mean{num_nodes.float().mean():.1f}) print(f边数: min{num_edges.min()}, max{num_edges.max()}, mean{num_edges.float().mean():.1f})这里打印节点数和边数的统计意义在于提前暴露脏数据QM9 中的分子最多九个重原子加氢后总原子数一般不超过 23如果你看到某个样本的节点数异常大大概率是下载或解析出了问题而不是真实分子。另一个需要核对的字段是y的数值范围QM9 的能量单位是 Hartree值在几百到一千多之间数值尺度很大后面训练时必须做标准化这个问题在第 5 章会细讲。数据切分也是在这个环节就定下来的。我一般不用随机切分而是先把数据集按分子中的碳原子数排序再依次划分训练、验证、测试集比例 8:1:1。这样保证每组数据都有相似的分子尺寸分布避免出现训练集全是小分子、测试集全是大分子的尴尬局面。4. 从消息传递到能量回归模型实现与训练循环4.1 模型结构一个可以直接复制的 MPNN项目里最重要的模型结构是消息传递神经网络Message Passing Neural Network。它的核心逻辑是每个原子的特征向量先聚合邻居原子的信息更新后再进入下一层经过 2 到 3 层消息传递后每个节点向量都包含了以它为中心、半径若干键之内的化学环境最后用求和池化把所有原子向量聚合成整分子表示接一个多层感知机输出能量。import torch import torch.nn.functional as F from torch import nn from torch_geometric.nn import MessagePassing, global_add_pool class MPNLayer(MessagePassing): def __init__(self, hidden_dim, edge_dim): super().__init__(aggradd) self.edge_mlp nn.Sequential( nn.Linear(hidden_dim * 2 edge_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim), ) self.residual nn.Linear(hidden_dim, hidden_dim) def forward(self, x, edge_index, edge_attr): out self.propagate(edge_index, xx, edge_attredge_attr) return F.silu(out self.residual(x)) def message(self, x_j, x_i, edge_attr): # x_j 是邻居节点x_i 是中心节点 msg_in torch.cat([x_i, x_j, edge_attr], dim-1) return self.edge_mlp(msg_in) class EnergyGNN(nn.Module): def __init__(self, node_dim, edge_dim, hidden_dim128, num_layers3): super().__init__() self.node_embed nn.Linear(node_dim, hidden_dim) self.layers nn.ModuleList() for _ in range(num_layers): self.layers.append(MPNLayer(hidden_dim, edge_dim)) self.readout nn.Sequential( nn.Linear(hidden_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1), ) def forward(self, data): x self.node_embed(data.x) for layer in self.layers: x layer(x, data.edge_index, data.edge_attr) x global_add_pool(x, data.batch) return self.readout(x).squeeze(-1)这个模型有几个参数值得细说。第一num_layers我默认设成 3这是我在 QM9 上反复试出来的折中点。层数加到 6 以上验证误差反而会回升分子太小三层消息传递已经覆盖了原子周围 3 个键的局部环境再堆层数只会把不同原子的表征抹平。第二消息函数里同时使用x_i和x_jx_i是中心节点x_j是邻居节点拼接后再过 MLP这样模型可以学到“中心原子和邻居原子分别是什么、它们以什么键相连”的联合表示。第三残差连接必不可少这里跳过了消息聚合的直接相加缓解了深层网络的梯度消失与过平滑。global_add_pool在这里做的不是取平均而是求和这也是刻意的。能量是广延量分子越大原子越多总能量绝对数值越大求和池化保留了“原子数”信息平均池化会把它抹掉。如果你改成平均池化预测误差会显著上升项目里如果发现预测值偏小第一个就检查这里。4.2 训练循环与评估指标训练过程的关键不是模型代码本身而是 loss 和评估单位的选择。我推荐用均方误差作为训练损失、平均绝对误差作为评估指标最后统一换算成 kcal/mol 来对外汇报结果。下面是一份可以直接跑起来的训练脚本骨架import torch from torch_geometric.loader import DataLoader def make_loaders(dataset, batch_size128, val_ratio0.1, test_ratio0.1): num_total len(dataset) num_val int(num_total * val_ratio) num_test int(num_total * test_ratio) num_train num_total - num_val - num_test idx torch.randperm(num_total) train_idx idx[:num_train] val_idx idx[num_train:num_train num_val] test_idx idx[num_train num_val:] train_loader DataLoader([dataset[i] for i in train_idx], batch_sizebatch_size, shuffleTrue) val_loader DataLoader([dataset[i] for i in val_idx], batch_sizebatch_size, shuffleFalse) test_loader DataLoader([dataset[i] for i in test_idx], batch_sizebatch_size, shuffleFalse) return train_loader, val_loader, test_loader targets torch.cat([d.y[:, 0:1] for d in dataset], dim0) target_mean targets.mean() target_std targets.std() model EnergyGNN(node_dimdataset.num_node_features, edge_dim4) optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemin, factor0.5, patience5) train_loader, val_loader, test_loader make_loaders(dataset) for epoch in range(50): model.train() total_loss 0.0 for data in train_loader: normalized_target (data.y[:, 0] - target_mean) / target_std pred model(data) loss torch.nn.functional.mse_loss(pred, normalized_target) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * data.num_graphs model.eval() val_preds, val_targets [], [] with torch.no_grad(): for data in val_loader: pred model(data) * target_std target_mean val_preds.append(pred) val_targets.append(data.y[:, 0]) val_preds torch.cat(val_preds) val_targets torch.cat(val_targets) mae_eV (val_preds - val_targets).abs().mean().item() mae_kcal mae_eV * 23.0605 # 1 eV 约等于 qm9/kcal 换算系数这里做单位换算 scheduler.step(mae_eV) if (epoch 1) % 5 0: print(fEpoch {epoch1:02d} | loss {total_loss/len(train_loader):.4f} | val MAE {mae_eV:.3f} eV)训练部分有三个我踩过的坑直接写在这里。第一优化器用 Adam 时学习率 1e-3 是安全起点但如果 loss 在训练开始就不降或者剧烈震荡把学习率降到 5e-4 比换优化器更快见效。第二ReduceLROnPlateau的patience设为 5意味着连续 5 个 epoch 没有下降才减半学习率给模型足够的搜索空间不要一看到验证误差平台期就手动调参。第三单位换算是报告结果时最容易乱的地方QM9 原始标签单位是 Hartreetarget_std和target_mean必须在训练集上计算验证集和测试集只能复用这两个数绝对不要在测试集上重新计算否则等于把测试集信息泄进评估流程。5. 分子图构建与训练中的 5 个典型坑5.1 去掉氢原子导致误差翻倍现象模型结构没变只把Chem.AddHs去掉验证集 MAE 立刻从几十 meV 涨到百 meV 级别。原因QM9 中 U0 能量是包含全原子贡献的量子化学结果。每个分子的 C-H 键数量不同氢原子及其与碳骨架的相对位置直接影响体系的零点能与振动贡献。把氢原子从图里删掉等于人为抹掉了一部分能量信息源模型只能靠碳骨架间的差异去猜误差必然变大。解决保留显式氢原子。如果你因为训练速度原因想删氢至少把氢原子总数作为全局特征拼到global_add_pool的输出后再进读出网络。我对比过两个方案的验证误差保留氢原子的建模方式误差明显更低。5.2 原子序数 one-hot 越界现象训练中途报RuntimeError: index out of bounds而且报错位置在模型第一层而非数据层。原因ATOM_ORDER只覆盖五个元素但 RDKit 解析时可能遇到其他元素或者数据集内部混入了原子序数为 0 的占位符。one-hot 向量的长度不够时越界的原子序数会让索引直接超出矩阵维度。解决在所有 one-hot 映射代码里加一个兜底维度“其他”同时在外层做一次原子种类统计打印出现过的原子序数集合。我习惯在smiles_to_graph里加一个unknown_count计数器统计“其他”类原子的数量一旦超过零就要怀疑数据源是否跑偏了。5.3 数据随机切分造成同骨架泄漏现象训练误差和验证误差都很好看但把模型放到外部测试集上时误差明显变大甚至训练过程中验证误差就开始回升。原因随机切分时分子骨架相似的异构体可能同时出现在训练集和验证集模型在训练阶段“见过”这些骨架的局部模式验证误差被严重低估。解决切分前先按 Bemis-Murcko 骨架聚类让同一骨架的分子尽量分到同一组。我用 RDKit 的MurckoScaffoldGeneric算骨架再以小骨架为单位做分层划分。这个改动会让验证误差数值变大但评估结果更可信。不要为了追求好看的验证数字而用随机切分最终坑的是你自己。5.4 目标值未标准化导致收敛缓慢现象训练 loss 在 1e-3 量级震荡两三个小时没有明显下降学习率调整后要么不收敛要么直接发散。原因QM9 的 U0 数值范围很大直接把原始值当回归目标MSE 的梯度尺度会被大数值目标放大到不可控模型优化的重心全都偏到了对齐中心点而不是学习局部差异。解决在训练循环里对目标做零均值单位方差标准化推理时反标准化回真实单位。这一段逻辑在 4.2 节代码里已经实现重点是把target_mean和target_std只在训练集上计算并且用 Python 浮点数而非 Tensor 存下来避免加载新数据时维度出问题。5.5 消息传递层数过深反而过平滑现象层数从 3 加到 6验证 MAE 不降反升把每层节点向量打印出来做可视化发现不同分子的节点表征趋于相似。原因多次消息传递让每个原子的表示扩散到整个分子局部化学环境被全局信息淹没模型失去了区分官能团的能力。分子图本身很小过大的感受野没有价值。解决层数控制在 2 到 4 之间默认 3必须在残差连接存在的前提下再谈加层。如果你预测的是超大分子可以考虑用跳跃连接把不同层级的节点表示拼起来而不是纯粹加深。6. 用误差分布判断模型是否真的学会了化学训练完成不是终点真正的验收要看误差落在哪些分子上。我会做两类分析一是按原子数分组统计误差二是按分子骨架官能团类型分组看系统偏差。前者帮我们发现样本量不足的区域后者帮我们理解模型是否抓住了真正的化学规律。import pandas as pd import torch def analyze_error_by_size(model, loader, target_mean, target_std): model.eval() rows [] with torch.no_grad(): for data in loader: pred model(data) * target_std target_mean target data.y[:, 0] num_nodes data.num_nodes batch data.batch for i in range(data.num_graphs): row { mae: abs(pred[i].item() - target[i].item()), num_nodes: (batch i).sum().item() } rows.append(row) df pd.DataFrame(rows) stats df.groupby(num_nodes)[mae].agg([mean, count]) return stats这个函数输出的表格能回答两个问题模型是否在小分子上更好、大数据量的分子段是否误差更低。如果某个原子数段的分子在训练集中出现次数很少但误差异常高说明不是模型问题而是数据覆盖不足从建模角度直接用这个结论和上级沟通数据扩充方案比盲目调参有说服力得多。官能团级别的分析做起来也不复杂。对每个测试集分子用 RDKit 提取是否含羧基、氨基、醛基等然后分别统计预测误差。如果模型对含某个官能团的分子系统性偏高 20 meV说明当前消息传递结构对电负性原子的局部环境建模不够这时候最有效的改法是把键长加入edge_attr和pos坐标并用距离加权图来替换纯拓扑图而不是继续堆层数。我在这套系统上试过距离加权方案收敛速度没有变快但最终误差又能往下压一档代价是按欧氏距离构造边时需要设定截断半径通常取 3 埃到 4 埃之间太短会漏掉弱相互作用太长会把无关原子强行连上边。最后一件事是相信数据胜过相信直觉。每次我做完一次调整都会把新模型和旧模型在相同测试集上的逐样本误差画成散点图事先写好误差绝对值降低的统计比例再验收防止“感觉准了”的错觉。如果你的项目数据不是 QM9 而是自己的分子库同样的图形化对比流程也能直接复刻只需要改数据加载部分模型与训练脚本原样可用。希望这套从数据构建到误差拆解的经历能帮到你少走我走过的弯路。本文还有配套的精品资源点击获取
