简介基于图神经网络的分子能量预测是深度学习在化学领域的重要应用这份资源面向机器学习与化学交叉方向的研究者和学习者提供Python完整源码与配套数据包覆盖从分子图表示到能量回归预测的实现流程。压缩包共27个文件包含8个Python脚本、7个CSV数据文件、3个模型权重文件、文本说明、分子结构文件及结果图片等整体大小仅5.4MB下载和复现都很便捷。目前已有509人学习下载适合用于课程设计、科研入门或材料/药物筛选相关项目的算法验证。值得重点关注的是代码覆盖数据预处理、分子图构建、GNN模型定义、训练评估与结果可视化等完整环节配合CSV数据与模型权重可快速跑通实验还能结合rdkit、PyTorch Geometric等工具链深入理解原子特征聚合、邻居信息传递与能量预测的建模思路。1. GNN 分子能量预测完整源码和数据包都到手先别急着训练基于 GNN 的分子能量预测Python 完整源码和数据包都放到面前时最错误的打开方式就是找到 B.py 直接训练。我拆过这套 Molecular-energy-prediction-with-GNN 的资源包第一轮就栽在文件结构上raw、processed、A/B 两套流程、多个带数字后缀的脚本判断哪个入口才是主线比调模型参数更费时间。分子能量预测要解决的实际问题很具体给定 SMILES 或 mol 文件预测 U0、H298、G 这类量子化学能量服务于药物筛序和材料设计的早期排序。适合两类人刚入门图神经网络、想跑通一个真实化学数据集的从业者以及手里有分子数据但不知道怎么转化成图结构的化学信息学工程师。2. 从 SMILES 到分子图把数据包里的文件先读明白再动手拿到压缩包先别急着跑代码。GNN 的整个技术栈里最不值得纠结的其实是模型消息传递层就那么几种写法真正决定项目成败的是分子表示和数据组织方式。同一个 SMILES加了氢和没加氢在 GNN 看来是两个完全不同的图同一批数据随机划分和按骨架划分测试集 MAE 能差一个数量级。所以这一章先从文件结构讲起再给出分子转图的完整脚本最后落到 PyG 的 Dataset 组织上。2.1 压缩包里这些文件分别是什么角色先看文件清单这套资源里几个关键文件和目录的作用如下文件/目录在项目里的角色raw/原始分子文件目录常见 .mol 和 .xmlprocessed/预处理后的图缓存PyG Dataset 自动读写smiles_CAS.txt / cas20210902.txtSMILES 与 CAS 号对照表substance-693027.mol / mol_2D.mol单分子 2D 结构文件reaction-10124203.xml反应或标注文件含分子关系信息mol_gnn.py模型定义主文件B.py / BB.py / A.py / A_loder.py / A_write.py数据处理与训练脚本test.py / test21.py / test1.txt测试脚本与单一分子测试样本要特别注意 A_ 和 B_ 两组脚本。A_loder.py、A_write.py 看名字像数据读写流程B.py、BB.py 更像训练主线。读这种多脚本项目我一般先用 grep 找哪个脚本 import 了 mol_gnn.py顺着调用关系判断主线而不是每个脚本都跑一遍。test21.py 这种带数字的脚本大概率是某个阶段跑单一分子或某批 CAS 样本的遗留测试不一定要纳入主流程。raw 和 processed 是 PyG Dataset 的标配目录。raw 放原始文件processed 放处理好的图对象缓存。第二次运行 Dataset 时只要 processed 目录里已经有文件PyG 会直接跳过 process()这个设计在反复调参时能省掉大量重复转换时间。如果你改了分子特征记得手动删掉 processed 下的缓存否则跑的永远是旧数据。2.2 用 rdkit 把 SMILES 转成分子图节点特征和边特征怎么定分子转图的核心是把原子、化学键、局部环境变成张量。下面这段函数是这套流程里最基础的一环from rdkit import Chem import torch def smiles_to_graph(smiles: str, add_hs: bool True) - dict: 把 SMILES 字符串转成 GNN 需要的 x / edge_index / edge_attr mol Chem.MolFromSmiles(smiles) if mol is None: raise ValueError(f无法解析的 SMILES: {smiles}) if add_hs: # 显式加氢节点数会明显增加训练和推理必须保持一致 mol Chem.AddHs(mol) atom_feats [] for atom in mol.GetAtoms(): atom_feats.append([ float(atom.GetAtomicNum()), # 原子序数区分元素种类 float(atom.GetTotalNumHs()), # 连接氢数反映杂化状态 float(atom.GetFormalCharge()), # 形式电荷 1.0 if atom.IsInRing() else 0.0, # 是否成环 ]) x torch.tensor(atom_feats, dtypetorch.float) edge_index [] # shape (2, E)每条有向边两个端点 edge_attr [] # 与 edge_index 一一对应的边特征 for bond in mol.GetBonds(): u, v bond.GetBeginAtomIdx(), bond.GetEndAtomIdx() bond_feats [ float(bond.GetBondType()), # 键型枚举单键/双键/三键 1.0 if bond.GetIsAromatic() else 0.0 # 是否芳香键 ] # 化学键是无向的消息传递里一般展开成两条有向边 edge_index.append([u, v]); edge_attr.append(bond_feats) edge_index.append([v, u]); edge_attr.append(bond_feats) edge_index torch.tensor(edge_index, dtypetorch.long).t().contiguous() edge_attr torch.tensor(edge_attr, dtypetorch.float) return {x: x, edge_index: edge_index, edge_attr: edge_attr}这段脚本解决的是“分子怎么变成张量”的问题。x 是节点特征矩阵每行对应一个原子edge_index 是图的邻接关系PyG 要求 shape 是 (2, 边数)所以代码里用 t() 转置edge_attr 和 edge_index 对齐每条边带自己的特征。整个函数输出可以直接传给 torch_geometric.data.Data。参数说明里最需要注意的是 add_hsTrue。显式加氢后乙醇 C2H6O 的图从 3 个重原子变成 9 个节点模型输入维度没变但图结构完全不同。训练时加了氢、推理时没加模型精度会立刻崩掉甚至维度都对不上。节点特征里 GetTotalNumHs 和 IsInRing 是 RDKit 中区分相似原子环境最便宜的组合比如苯环上的碳和脂肪链上的碳原子序数相同但环标志和氢数能把它们分开。后续想加更丰富的特征可以把原子序数改成可学习的 embedding但先跑通流程不需要。2.3 把数据集组织成 PyG 的 Datasetprocessed 缓存到底存了什么单分子转成图之后下一步是把整个数据集组织成 PyG 能直接喂给 DataLoader 的对象。标准做法是继承 torch_geometric.data.Datasetfrom torch_geometric.data import Dataset, Data import os import torch class MolEnergyDataset(Dataset): def __init__(self, root, smiles_path, energy_path): self.smiles read_smiles(smiles_path) # 自己按数据表格式实现 self.energy read_energy(energy_path) # 返回与 smiles 一一对应的能量列表 super().__init__(root) property def processed_file_names(self): return [mol_energy.pt] # 存在 processed/ 下的缓存文件名 def process(self): data_list [] for s, e in zip(self.smiles, self.energy): graph smiles_to_graph(s, add_hsTrue) data Data(xgraph[x], edge_indexgraph[edge_index], edge_attrgraph[edge_attr], ytorch.tensor([e], dtypetorch.float)) data_list.append(data) # 把所有图打包保存第二次运行直接读缓存跳过 process() torch.save(data_list, os.path.join(self.processed_dir, mol_energy.pt)) def len(self): return len(self.smiles) def get(self, idx): return torch.load(os.path.join(self.processed_dir, mol_energy.pt))[idx]这个 Dataset 类把上一小节的转图函数接进 PyG 的标准做法。process() 负责把所有 SMILES 批量转成 Data 对象并缓存第二次实例化时PyG 发现 processed 目录里已有 mol_energy.pt会自动跳过 process()。get() 按索引取出单个样本DataLoader 会自动在此基础上聚合成 batch。root 是数据根目录PyG 会在 root 下自动建 raw/ 和 processed/ 两个子目录原包里那两目录就是这么来的。我没有在init里传 transform因为分子图不需要随机裁剪这类增强如果你的数据里混入大量同分异构体反而应该在划分阶段处理而不是在 transform 里做扰动。energy 的读取函数要和手里的数据表字段对齐常见的是 CSV 里一列 SMILES、一列能量值单位最好先在 CSV 里统一成 eV后面训练和评估都少一层换算。3. 构建 GNN 模型消息传递层、读出函数与推理脚本怎么串起来数据准备到位后才轮到模型。GNN 做分子能量预测的套路很固定消息传递层提取局部化学环境读出函数把节点特征聚合为整图特征最后接一个回归头输出能量。这一章会先讲清楚为什么这么选型再给出能直接跑的模型代码最后讲怎么快速读懂包里的 mol_gnn.py 和 test.py。3.1 为什么消息传递框架适合能量预测GCN 和 GIN 差在哪GNN 处理分子的核心假设是原子的化学环境决定它的局部能量贡献而化学环境可以靠“这个原子 它的邻居 邻居的邻居”来描述。消息传递层做的事情就是把邻居特征反复聚合到中心节点上经过 k 层后每个节点包含 k 跳范围内的结构信息。分子能量是整张图的性质所以最后要用一个读出函数把全部节点聚合成图级向量。GCN 和 GIN 的差别在于聚合方式。GCN 默认做归一化后的均值聚合GIN 做加和聚合并给中心节点加一个可学习的权重。对分子这种小图均值会稀释掉原子的个体信息而 GIN 在理论上接近 WL test 的同构判别能力能更好地区分结构相似的分子。经验上分子能量预测用 GINConv 或 GINEConv 比 GCN 稳这是第一个可以直接照抄的选型结论。3.2 一个能跑通的能量回归模型模型代码不长但几个关键细节会影响最终精度import torch import torch.nn.functional as F from torch_geometric.nn import GINConv, global_add_pool class GNNMolEnergy(torch.nn.Module): def __init__(self, in_dim4, hidden_dim128, out_dim1, num_layers3): super().__init__() self.convs torch.nn.ModuleList() for i in range(num_layers): # GINConv 的 MLP 里放 BatchNorm 和 ReLU比线性层直出稳定 self.convs.append(GINConv(torch.nn.Sequential( torch.nn.Linear(in_dim if i 0 else hidden_dim, hidden_dim), torch.nn.BatchNorm1d(hidden_dim), torch.nn.ReLU(), torch.nn.Linear(hidden_dim, hidden_dim), torch.nn.ReLU(), ))) self.readout global_add_pool # 加和读出能量近似原子贡献之和 self.regressor torch.nn.Sequential( torch.nn.Linear(hidden_dim, hidden_dim), torch.nn.ReLU(), torch.nn.Linear(hidden_dim, out_dim), ) def forward(self, x, edge_index, batch): for conv in self.convs: x conv(x, edge_index) x self.readout(x, batch) # 聚合到图级 return self.regressor(x)网络结构分三段消息传递、读出、回归。forward 里 conv 循环做邻居聚合每层输出维度保持 hidden_dimglobal_add_pool 把同一个 batch 里所有节点按 batch 向量分组求和输出 shape 是 (batch_size, hidden_dim)最后的 regressor 把图级向量压到 1 维能量。整个模型输入是 DataLoader 吐出的 batchbatch 向量由 PyG 自动生成。参数上hidden_dim128 是通用起点小数据集可以降到 64。num_layers3 是分子图的常规选择不要一味加深后面避坑章会专门讲。BatchNorm1d 在分子图上有个坑不同 batch 包含的原子数差别大BN 统计量会漂移如果发现验证集 loss 震荡可以把 BN 换成 LayerNorm。edge_attr 在这个版本里没用上想用键类型参与聚合就把 GINConv 换成 GINEConv并把 edge_attr 传进 conv。提示这个模型学的是分子图拓扑到能量的映射不适用于需要三维构象精确能量的场景比如反应路径计算。3.3 推理脚本怎么读从 mol_gnn.py 到 test.py 的调用关系这套包里 mol_gnn.py 是模型定义主文件test.py 和 test21.py 是测试入口。读多脚本项目时先用 grep 找出哪个脚本 import 了 mol_gnn.py顺着调用关系看主线。常见结构是 mol_gnn.py 只放模型test.py 负责加载 checkpoint 并评估test21.py 是针对单一分子或某批 CAS 样本的回归测试。推理时最容易漏的是 batch 向量的构造单分子也要走一遍import torch from torch_geometric.data import Data # 假设 checkpoint 是训练完成后保存的状态字典 model GNNMolEnergy(in_dim4, hidden_dim128, out_dim1, num_layers3) ckpt torch.load(model_epoch100.pt, map_locationcpu) model.load_state_dict(ckpt[model_state_dict] if model_state_dict in ckpt else ckpt) model.eval() with torch.no_grad(): graph smiles_to_graph(CCO, add_hsTrue) # 乙醇 data Data(**graph) # 单分子也要构造 batch 向量全 0 即可 batch torch.zeros(data.x.size(0), dtypetorch.long) energy model(data.x, data.edge_index, batch) print(predicted energy:, energy.item())这段代码演示了模型训练完之后的原始调用过程。注意单分子的 batch 向量是一个全零的长整数张量表示所有节点属于同一个图。如果推理时忘了这一步GNN 会把整个分子当成一个样本维度和训练时完全对不上。4. 训练与评估误差压到多少才算能用模型能跑通只是开始分子能量预测里真正影响结果的是训练配置和数据划分。这一章解决三个问题损失函数选什么、训练循环怎么写、数据怎么划分才能反映真实泛化能力。4.1 损失函数和评估指标MAE、RMSE 与单位陷阱回归任务首选 MAE因为它和能量单位绑定能直接判断“预测值平均偏了多少”。RMSE 对大误差分子更敏感如果想暴露特别离谱的样本看 RMSE。R² 适合汇报给非技术同事但分子能量预测里 R² 很容易被数据方差撑得很好看0.98 不代表 MAE 小。单位问题比指标选择更致命。不同来源的数据能量单位可能完全不同单位换算关系常见场景eV1 eV 23.06 kcal/molQM9 的能量、HOMO/LUMOHartree1 Hartree 27.2114 eVGaussian 等量子化学软件输出kcal/mol原子化能、热力学性质化学手册、文献数据注意如果你的源数据是 eV 而模型输出不加激活训练初期 loss 下降不明显是正常的先跑一个 epoch 检查预测数值范围再判断。在 QM9 上3 层 GIN、128 维隐藏层验证集 MAE 在 0.05~0.15 eV 属于正常区间如果高于 0.3 eV先怀疑数据没洗干净而不是模型不够强。这个经验值能帮你判断一次训练到底有没有意义。4.2 一个标准的训练循环L1 损失、Adam 与学习率衰减训练代码是 PyG 任务的标准模板下面这段可以直接套用import torch import torch.optim as optim from torch_geometric.loader import DataLoader model GNNMolEnergy(in_dim4, hidden_dim128, out_dim1) optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-5) scheduler optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemin, factor0.5, patience10) loss_fn torch.nn.L1Loss() train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) val_loader DataLoader(val_dataset, batch_size128, shuffleFalse) def evaluate(loader): model.eval() errs [] with torch.no_grad(): for batch in loader: pred model(batch.x, batch.edge_index, batch.batch).squeeze(-1) errs.append((pred - batch.y).abs()) return torch.cat(errs).mean().item() for epoch in range(300): model.train() total 0.0 for batch in train_loader: optimizer.zero_grad() pred model(batch.x, batch.edge_index, batch.batch).squeeze(-1) loss loss_fn(pred, batch.y) loss.backward() optimizer.step() total loss.item() * batch.num_graphs val_mae evaluate(val_loader) scheduler.step(val_mae) # 验证集 MAE 不再降时减半学习率 if epoch % 10 0: print(fepoch {epoch:3d} train_loss{total / len(train_dataset):.4f} val_mae{val_mae:.4f})batch_size64 对分子图来说算紧凑分子平均 20~40 个节点显存占用不大。如果数据集里混着大分子要按原子数分桶后面避坑章会讲。L1Loss 对异常值不敏感符合能量预测的诉求。scheduler 用验证集 MAE 驱动patience10 表示 10 个 epoch 不下降就减半学习率300 个 epoch 足够收敛。训练时我只盯两个信号验证集 MAE 是不是在稳步下降学习率减半后有没有再跳一截。训练集 loss 降到接近 0 不是好事说明模型开始背数据。4.3 数据划分随机划分最容易高估模型这是没有后悔药的坑分子数据划分不能直接 random split原因很直接同一骨架的分子结构高度相似随机划分会把大量近亲分子同时分到训练集和验证集。模型看到的是结构联想不是真正的泛化。更严格的做法是按 Murcko 骨架分桶from rdkit import Chem from rdkit.Chem.Scaffolds import MurckoScaffold def scaffold_split(smiles_list, valid_ratio0.2): scaffolds {} for idx, s in enumerate(smiles_list): mol Chem.MolFromSmiles(s) if mol is None: continue # 取 Murcko 骨架作为桶 key同一骨架的分子必须进同一份数据 scaf MurckoScaffold.MurckoScaffoldSmiles(molmol) scaffolds.setdefault(scaf, []).append(idx) train_idx, val_idx [], [] for scaf, idxs in sorted(scaffolds.items(), keylambda kv: len(kv[1]), reverseTrue): if len(val_idx) / max(len(smiles_list), 1) valid_ratio: val_idx idxs else: train_idx idxs return train_idx, val_idx这个划分策略把分子按 Murcko 骨架分类再把整个骨架桶分给验证集保证验证集里的分子和训练集没有共享核心骨架。valid_ratio 控制验证集规模这里按分子数近似控制。更严格的做法是先按 CAS 号去重再按骨架分桶防止同一个化合物在数据表里出现两次。如果只是为了跑通流程随机划分可以接受但要汇报模型泛化能力必须说明用的是哪种划分。同样一个模型随机划分 MAE 0.06 eV骨架划分 0.15 eV都很正常。这个差距不是模型不行是评估的严格程度不同。5. 避坑分子能量预测最容易翻车的五个点这些坑我在拆源码和调参过程中都踩过每一条都是现象、原因、解决三层。前三个在数据侧后两个在模型侧。5.1 数据侧的三个坑单位、加氢、划分坑 1能量差了几个数量级loss 还降得挺好看现象验证集 MAE 稳定在 0.02但把预测值拿出来和文献能量一比差了几十倍。原因数据的能量单位不统一一部分是 eV一部分是 Hartree或者混在同一列里。GNN 的输出没有任何物理约束它会自动学会拟合“当前这一批数字”但换到另一批单位就全乱。解决训练前对目标值 y 做数值范围检查打印 min、max、std。Hartree 转 eV 乘以 27.2114kcal/mol 转 eV 除以 23.06。我一般把单位换算单独写一个函数并在 Dataset 加载时断言 y 的范围在合理区间超了就报错。坑 2同样的模型加氢和不加氢结果完全不一样现象训练时用 AddHs推理时忘了加氢模型输出严重偏离甚至特征维度对不上。原因氢原子作为节点参与消息传递图的节点数和边数都变了模型学到的分布完全被破坏。解决把 add_hs 作为 Dataset 和推理脚本的固定参数两者必须来自同一个入口。做一致性检查同一个 SMILES 在训练前和推理前的节点数必须一致否则直接抛异常。这个检查 10 秒钟能省掉一整晚的排错。坑 3随机划分让测试集 MAE 变得虚低现象交叉验证 MAE 0.06 eV投到一批新骨架分子上误差直接翻三倍。原因随机划分把同骨架的类似分子同时分到训练集和测试集模型是在做结构联想不是泛化。解决按 Murcko 骨架分桶划分或至少按 CAS 号去重后随机划分。报告精度时注明划分方式否则不同论文之间没法横向比较。5.2 模型侧的两个坑过平滑与坐标混用坑 4消息传递层加到 5 层精度反而下降现象层数从 2 加到 5训练 loss 能下降验证 MAE 先降后升所有节点的表示越来越相似。原因分子图直径小苯环体系 2~3 跳就能覆盖全局。层数继续加深等于重复做同样的聚合节点特征被过度平均这就是典型的过平滑现象。解决分子能量预测一般 2~4 层就够全局信息靠读出层而不是靠加深网络。真需要更大感受野用跳跃连接或 JK-net 在层间加权组合。坑 5把 mol_2D.mol 当 3D 构象用预测构象能量一塌糊涂现象想用这套 GNN 预测某个分子的力场单点能输入 2D 结构文件结果 MAE 高达几个 kcal/mol完全不可用。原因.mol 文件里是 2D 坐标没有键长键角信息。GNN 学的是拓扑到能量的映射不是三维构象到能量的映射本质上有适用边界。解决先明确预测目标。QM9 类数据给出的是分子图对应的电子能量图结构就够如果目标是构象能量或势能面需要用 ETKDG 或 RDKit 生成 3D 构象再考虑 3D GNN 或等变网络。拿到新数据时我通常会按这个顺序过一遍打印 y 的 min/max/std随机抽三个 SMILES 检查加氢前后节点数统计骨架桶数量决定划分方式固定 hidden_dim128、num_layers3 快速跑通最后用单样本推理核对输出范围。这套习惯能避免 80% 的无效训练。6. 进阶用多任务改一版一次预测四个能量再做误差归因前面几章把单能量预测跑通了但实际项目里往往不只有一个目标。QM9 同时提供 U0、U298、H298、G298、HOMO、LUMO 等多项性质你当然可以训练六个模型但数据复用效率太低。更合理的做法是在最后一个隐藏层后加一个多任务回归头共享底层消息传递网络class MultiTaskEnergy(torch.nn.Module): def __init__(self, hidden_dim128, n_tasks4): super().__init__() # 复用前面 GNNMolEnergy只把它当成编码器 self.encoder GNNMolEnergy(in_dim4, hidden_dimhidden_dim, out_dimhidden_dim, num_layers3) self.heads torch.nn.ModuleList([ torch.nn.Linear(hidden_dim, 1) for _ in range(n_tasks) ]) def forward(self, x, edge_index, batch): h self.encoder(x, edge_index, batch) # 共享分子表示 return torch.cat([head(h) for head in self.heads], dim-1)关键点是每个任务只有一个线性头消息传递部分是共用的。U0 和 H298 高度相关共享表示能让数据互相补充小数据集下比六个独立模型稳得多。训练时把 L1Loss 改成对所有任务输出取平均即可。模型训完之后下一步不是急着调参而是做误差归因。把验证集里误差最大的 20 个分子挑出来按 Murcko 骨架聚类看模型到底对哪类结构无能为力。我见过最典型的场景是误差集中在含氮杂环体系上说明训练数据里这类结构太少需要回数据层面补样本而不是继续调学习率。这种黑匣子分析比看 loss 曲线有用得多。从那以后我每次拿到新的分子能量数据都强制走一遍完整流程确认单位固定加氢逻辑按骨架划分用一个 3 层 GIN 快速跑通再做多任务和误差归因。这套流程帮我省掉了大量试错时间也希望帮到你。本文还有配套的精品资源点击获取
GNN分子能量预测实战：从SMILES到分子图构建与模型训练
简介基于图神经网络的分子能量预测是深度学习在化学领域的重要应用这份资源面向机器学习与化学交叉方向的研究者和学习者提供Python完整源码与配套数据包覆盖从分子图表示到能量回归预测的实现流程。压缩包共27个文件包含8个Python脚本、7个CSV数据文件、3个模型权重文件、文本说明、分子结构文件及结果图片等整体大小仅5.4MB下载和复现都很便捷。目前已有509人学习下载适合用于课程设计、科研入门或材料/药物筛选相关项目的算法验证。值得重点关注的是代码覆盖数据预处理、分子图构建、GNN模型定义、训练评估与结果可视化等完整环节配合CSV数据与模型权重可快速跑通实验还能结合rdkit、PyTorch Geometric等工具链深入理解原子特征聚合、邻居信息传递与能量预测的建模思路。1. GNN 分子能量预测完整源码和数据包都到手先别急着训练基于 GNN 的分子能量预测Python 完整源码和数据包都放到面前时最错误的打开方式就是找到 B.py 直接训练。我拆过这套 Molecular-energy-prediction-with-GNN 的资源包第一轮就栽在文件结构上raw、processed、A/B 两套流程、多个带数字后缀的脚本判断哪个入口才是主线比调模型参数更费时间。分子能量预测要解决的实际问题很具体给定 SMILES 或 mol 文件预测 U0、H298、G 这类量子化学能量服务于药物筛序和材料设计的早期排序。适合两类人刚入门图神经网络、想跑通一个真实化学数据集的从业者以及手里有分子数据但不知道怎么转化成图结构的化学信息学工程师。2. 从 SMILES 到分子图把数据包里的文件先读明白再动手拿到压缩包先别急着跑代码。GNN 的整个技术栈里最不值得纠结的其实是模型消息传递层就那么几种写法真正决定项目成败的是分子表示和数据组织方式。同一个 SMILES加了氢和没加氢在 GNN 看来是两个完全不同的图同一批数据随机划分和按骨架划分测试集 MAE 能差一个数量级。所以这一章先从文件结构讲起再给出分子转图的完整脚本最后落到 PyG 的 Dataset 组织上。2.1 压缩包里这些文件分别是什么角色先看文件清单这套资源里几个关键文件和目录的作用如下文件/目录在项目里的角色raw/原始分子文件目录常见 .mol 和 .xmlprocessed/预处理后的图缓存PyG Dataset 自动读写smiles_CAS.txt / cas20210902.txtSMILES 与 CAS 号对照表substance-693027.mol / mol_2D.mol单分子 2D 结构文件reaction-10124203.xml反应或标注文件含分子关系信息mol_gnn.py模型定义主文件B.py / BB.py / A.py / A_loder.py / A_write.py数据处理与训练脚本test.py / test21.py / test1.txt测试脚本与单一分子测试样本要特别注意 A_ 和 B_ 两组脚本。A_loder.py、A_write.py 看名字像数据读写流程B.py、BB.py 更像训练主线。读这种多脚本项目我一般先用 grep 找哪个脚本 import 了 mol_gnn.py顺着调用关系判断主线而不是每个脚本都跑一遍。test21.py 这种带数字的脚本大概率是某个阶段跑单一分子或某批 CAS 样本的遗留测试不一定要纳入主流程。raw 和 processed 是 PyG Dataset 的标配目录。raw 放原始文件processed 放处理好的图对象缓存。第二次运行 Dataset 时只要 processed 目录里已经有文件PyG 会直接跳过 process()这个设计在反复调参时能省掉大量重复转换时间。如果你改了分子特征记得手动删掉 processed 下的缓存否则跑的永远是旧数据。2.2 用 rdkit 把 SMILES 转成分子图节点特征和边特征怎么定分子转图的核心是把原子、化学键、局部环境变成张量。下面这段函数是这套流程里最基础的一环from rdkit import Chem import torch def smiles_to_graph(smiles: str, add_hs: bool True) - dict: 把 SMILES 字符串转成 GNN 需要的 x / edge_index / edge_attr mol Chem.MolFromSmiles(smiles) if mol is None: raise ValueError(f无法解析的 SMILES: {smiles}) if add_hs: # 显式加氢节点数会明显增加训练和推理必须保持一致 mol Chem.AddHs(mol) atom_feats [] for atom in mol.GetAtoms(): atom_feats.append([ float(atom.GetAtomicNum()), # 原子序数区分元素种类 float(atom.GetTotalNumHs()), # 连接氢数反映杂化状态 float(atom.GetFormalCharge()), # 形式电荷 1.0 if atom.IsInRing() else 0.0, # 是否成环 ]) x torch.tensor(atom_feats, dtypetorch.float) edge_index [] # shape (2, E)每条有向边两个端点 edge_attr [] # 与 edge_index 一一对应的边特征 for bond in mol.GetBonds(): u, v bond.GetBeginAtomIdx(), bond.GetEndAtomIdx() bond_feats [ float(bond.GetBondType()), # 键型枚举单键/双键/三键 1.0 if bond.GetIsAromatic() else 0.0 # 是否芳香键 ] # 化学键是无向的消息传递里一般展开成两条有向边 edge_index.append([u, v]); edge_attr.append(bond_feats) edge_index.append([v, u]); edge_attr.append(bond_feats) edge_index torch.tensor(edge_index, dtypetorch.long).t().contiguous() edge_attr torch.tensor(edge_attr, dtypetorch.float) return {x: x, edge_index: edge_index, edge_attr: edge_attr}这段脚本解决的是“分子怎么变成张量”的问题。x 是节点特征矩阵每行对应一个原子edge_index 是图的邻接关系PyG 要求 shape 是 (2, 边数)所以代码里用 t() 转置edge_attr 和 edge_index 对齐每条边带自己的特征。整个函数输出可以直接传给 torch_geometric.data.Data。参数说明里最需要注意的是 add_hsTrue。显式加氢后乙醇 C2H6O 的图从 3 个重原子变成 9 个节点模型输入维度没变但图结构完全不同。训练时加了氢、推理时没加模型精度会立刻崩掉甚至维度都对不上。节点特征里 GetTotalNumHs 和 IsInRing 是 RDKit 中区分相似原子环境最便宜的组合比如苯环上的碳和脂肪链上的碳原子序数相同但环标志和氢数能把它们分开。后续想加更丰富的特征可以把原子序数改成可学习的 embedding但先跑通流程不需要。2.3 把数据集组织成 PyG 的 Datasetprocessed 缓存到底存了什么单分子转成图之后下一步是把整个数据集组织成 PyG 能直接喂给 DataLoader 的对象。标准做法是继承 torch_geometric.data.Datasetfrom torch_geometric.data import Dataset, Data import os import torch class MolEnergyDataset(Dataset): def __init__(self, root, smiles_path, energy_path): self.smiles read_smiles(smiles_path) # 自己按数据表格式实现 self.energy read_energy(energy_path) # 返回与 smiles 一一对应的能量列表 super().__init__(root) property def processed_file_names(self): return [mol_energy.pt] # 存在 processed/ 下的缓存文件名 def process(self): data_list [] for s, e in zip(self.smiles, self.energy): graph smiles_to_graph(s, add_hsTrue) data Data(xgraph[x], edge_indexgraph[edge_index], edge_attrgraph[edge_attr], ytorch.tensor([e], dtypetorch.float)) data_list.append(data) # 把所有图打包保存第二次运行直接读缓存跳过 process() torch.save(data_list, os.path.join(self.processed_dir, mol_energy.pt)) def len(self): return len(self.smiles) def get(self, idx): return torch.load(os.path.join(self.processed_dir, mol_energy.pt))[idx]这个 Dataset 类把上一小节的转图函数接进 PyG 的标准做法。process() 负责把所有 SMILES 批量转成 Data 对象并缓存第二次实例化时PyG 发现 processed 目录里已有 mol_energy.pt会自动跳过 process()。get() 按索引取出单个样本DataLoader 会自动在此基础上聚合成 batch。root 是数据根目录PyG 会在 root 下自动建 raw/ 和 processed/ 两个子目录原包里那两目录就是这么来的。我没有在init里传 transform因为分子图不需要随机裁剪这类增强如果你的数据里混入大量同分异构体反而应该在划分阶段处理而不是在 transform 里做扰动。energy 的读取函数要和手里的数据表字段对齐常见的是 CSV 里一列 SMILES、一列能量值单位最好先在 CSV 里统一成 eV后面训练和评估都少一层换算。3. 构建 GNN 模型消息传递层、读出函数与推理脚本怎么串起来数据准备到位后才轮到模型。GNN 做分子能量预测的套路很固定消息传递层提取局部化学环境读出函数把节点特征聚合为整图特征最后接一个回归头输出能量。这一章会先讲清楚为什么这么选型再给出能直接跑的模型代码最后讲怎么快速读懂包里的 mol_gnn.py 和 test.py。3.1 为什么消息传递框架适合能量预测GCN 和 GIN 差在哪GNN 处理分子的核心假设是原子的化学环境决定它的局部能量贡献而化学环境可以靠“这个原子 它的邻居 邻居的邻居”来描述。消息传递层做的事情就是把邻居特征反复聚合到中心节点上经过 k 层后每个节点包含 k 跳范围内的结构信息。分子能量是整张图的性质所以最后要用一个读出函数把全部节点聚合成图级向量。GCN 和 GIN 的差别在于聚合方式。GCN 默认做归一化后的均值聚合GIN 做加和聚合并给中心节点加一个可学习的权重。对分子这种小图均值会稀释掉原子的个体信息而 GIN 在理论上接近 WL test 的同构判别能力能更好地区分结构相似的分子。经验上分子能量预测用 GINConv 或 GINEConv 比 GCN 稳这是第一个可以直接照抄的选型结论。3.2 一个能跑通的能量回归模型模型代码不长但几个关键细节会影响最终精度import torch import torch.nn.functional as F from torch_geometric.nn import GINConv, global_add_pool class GNNMolEnergy(torch.nn.Module): def __init__(self, in_dim4, hidden_dim128, out_dim1, num_layers3): super().__init__() self.convs torch.nn.ModuleList() for i in range(num_layers): # GINConv 的 MLP 里放 BatchNorm 和 ReLU比线性层直出稳定 self.convs.append(GINConv(torch.nn.Sequential( torch.nn.Linear(in_dim if i 0 else hidden_dim, hidden_dim), torch.nn.BatchNorm1d(hidden_dim), torch.nn.ReLU(), torch.nn.Linear(hidden_dim, hidden_dim), torch.nn.ReLU(), ))) self.readout global_add_pool # 加和读出能量近似原子贡献之和 self.regressor torch.nn.Sequential( torch.nn.Linear(hidden_dim, hidden_dim), torch.nn.ReLU(), torch.nn.Linear(hidden_dim, out_dim), ) def forward(self, x, edge_index, batch): for conv in self.convs: x conv(x, edge_index) x self.readout(x, batch) # 聚合到图级 return self.regressor(x)网络结构分三段消息传递、读出、回归。forward 里 conv 循环做邻居聚合每层输出维度保持 hidden_dimglobal_add_pool 把同一个 batch 里所有节点按 batch 向量分组求和输出 shape 是 (batch_size, hidden_dim)最后的 regressor 把图级向量压到 1 维能量。整个模型输入是 DataLoader 吐出的 batchbatch 向量由 PyG 自动生成。参数上hidden_dim128 是通用起点小数据集可以降到 64。num_layers3 是分子图的常规选择不要一味加深后面避坑章会专门讲。BatchNorm1d 在分子图上有个坑不同 batch 包含的原子数差别大BN 统计量会漂移如果发现验证集 loss 震荡可以把 BN 换成 LayerNorm。edge_attr 在这个版本里没用上想用键类型参与聚合就把 GINConv 换成 GINEConv并把 edge_attr 传进 conv。提示这个模型学的是分子图拓扑到能量的映射不适用于需要三维构象精确能量的场景比如反应路径计算。3.3 推理脚本怎么读从 mol_gnn.py 到 test.py 的调用关系这套包里 mol_gnn.py 是模型定义主文件test.py 和 test21.py 是测试入口。读多脚本项目时先用 grep 找出哪个脚本 import 了 mol_gnn.py顺着调用关系看主线。常见结构是 mol_gnn.py 只放模型test.py 负责加载 checkpoint 并评估test21.py 是针对单一分子或某批 CAS 样本的回归测试。推理时最容易漏的是 batch 向量的构造单分子也要走一遍import torch from torch_geometric.data import Data # 假设 checkpoint 是训练完成后保存的状态字典 model GNNMolEnergy(in_dim4, hidden_dim128, out_dim1, num_layers3) ckpt torch.load(model_epoch100.pt, map_locationcpu) model.load_state_dict(ckpt[model_state_dict] if model_state_dict in ckpt else ckpt) model.eval() with torch.no_grad(): graph smiles_to_graph(CCO, add_hsTrue) # 乙醇 data Data(**graph) # 单分子也要构造 batch 向量全 0 即可 batch torch.zeros(data.x.size(0), dtypetorch.long) energy model(data.x, data.edge_index, batch) print(predicted energy:, energy.item())这段代码演示了模型训练完之后的原始调用过程。注意单分子的 batch 向量是一个全零的长整数张量表示所有节点属于同一个图。如果推理时忘了这一步GNN 会把整个分子当成一个样本维度和训练时完全对不上。4. 训练与评估误差压到多少才算能用模型能跑通只是开始分子能量预测里真正影响结果的是训练配置和数据划分。这一章解决三个问题损失函数选什么、训练循环怎么写、数据怎么划分才能反映真实泛化能力。4.1 损失函数和评估指标MAE、RMSE 与单位陷阱回归任务首选 MAE因为它和能量单位绑定能直接判断“预测值平均偏了多少”。RMSE 对大误差分子更敏感如果想暴露特别离谱的样本看 RMSE。R² 适合汇报给非技术同事但分子能量预测里 R² 很容易被数据方差撑得很好看0.98 不代表 MAE 小。单位问题比指标选择更致命。不同来源的数据能量单位可能完全不同单位换算关系常见场景eV1 eV 23.06 kcal/molQM9 的能量、HOMO/LUMOHartree1 Hartree 27.2114 eVGaussian 等量子化学软件输出kcal/mol原子化能、热力学性质化学手册、文献数据注意如果你的源数据是 eV 而模型输出不加激活训练初期 loss 下降不明显是正常的先跑一个 epoch 检查预测数值范围再判断。在 QM9 上3 层 GIN、128 维隐藏层验证集 MAE 在 0.05~0.15 eV 属于正常区间如果高于 0.3 eV先怀疑数据没洗干净而不是模型不够强。这个经验值能帮你判断一次训练到底有没有意义。4.2 一个标准的训练循环L1 损失、Adam 与学习率衰减训练代码是 PyG 任务的标准模板下面这段可以直接套用import torch import torch.optim as optim from torch_geometric.loader import DataLoader model GNNMolEnergy(in_dim4, hidden_dim128, out_dim1) optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-5) scheduler optim.lr_scheduler.ReduceLROnPlateau(optimizer, modemin, factor0.5, patience10) loss_fn torch.nn.L1Loss() train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) val_loader DataLoader(val_dataset, batch_size128, shuffleFalse) def evaluate(loader): model.eval() errs [] with torch.no_grad(): for batch in loader: pred model(batch.x, batch.edge_index, batch.batch).squeeze(-1) errs.append((pred - batch.y).abs()) return torch.cat(errs).mean().item() for epoch in range(300): model.train() total 0.0 for batch in train_loader: optimizer.zero_grad() pred model(batch.x, batch.edge_index, batch.batch).squeeze(-1) loss loss_fn(pred, batch.y) loss.backward() optimizer.step() total loss.item() * batch.num_graphs val_mae evaluate(val_loader) scheduler.step(val_mae) # 验证集 MAE 不再降时减半学习率 if epoch % 10 0: print(fepoch {epoch:3d} train_loss{total / len(train_dataset):.4f} val_mae{val_mae:.4f})batch_size64 对分子图来说算紧凑分子平均 20~40 个节点显存占用不大。如果数据集里混着大分子要按原子数分桶后面避坑章会讲。L1Loss 对异常值不敏感符合能量预测的诉求。scheduler 用验证集 MAE 驱动patience10 表示 10 个 epoch 不下降就减半学习率300 个 epoch 足够收敛。训练时我只盯两个信号验证集 MAE 是不是在稳步下降学习率减半后有没有再跳一截。训练集 loss 降到接近 0 不是好事说明模型开始背数据。4.3 数据划分随机划分最容易高估模型这是没有后悔药的坑分子数据划分不能直接 random split原因很直接同一骨架的分子结构高度相似随机划分会把大量近亲分子同时分到训练集和验证集。模型看到的是结构联想不是真正的泛化。更严格的做法是按 Murcko 骨架分桶from rdkit import Chem from rdkit.Chem.Scaffolds import MurckoScaffold def scaffold_split(smiles_list, valid_ratio0.2): scaffolds {} for idx, s in enumerate(smiles_list): mol Chem.MolFromSmiles(s) if mol is None: continue # 取 Murcko 骨架作为桶 key同一骨架的分子必须进同一份数据 scaf MurckoScaffold.MurckoScaffoldSmiles(molmol) scaffolds.setdefault(scaf, []).append(idx) train_idx, val_idx [], [] for scaf, idxs in sorted(scaffolds.items(), keylambda kv: len(kv[1]), reverseTrue): if len(val_idx) / max(len(smiles_list), 1) valid_ratio: val_idx idxs else: train_idx idxs return train_idx, val_idx这个划分策略把分子按 Murcko 骨架分类再把整个骨架桶分给验证集保证验证集里的分子和训练集没有共享核心骨架。valid_ratio 控制验证集规模这里按分子数近似控制。更严格的做法是先按 CAS 号去重再按骨架分桶防止同一个化合物在数据表里出现两次。如果只是为了跑通流程随机划分可以接受但要汇报模型泛化能力必须说明用的是哪种划分。同样一个模型随机划分 MAE 0.06 eV骨架划分 0.15 eV都很正常。这个差距不是模型不行是评估的严格程度不同。5. 避坑分子能量预测最容易翻车的五个点这些坑我在拆源码和调参过程中都踩过每一条都是现象、原因、解决三层。前三个在数据侧后两个在模型侧。5.1 数据侧的三个坑单位、加氢、划分坑 1能量差了几个数量级loss 还降得挺好看现象验证集 MAE 稳定在 0.02但把预测值拿出来和文献能量一比差了几十倍。原因数据的能量单位不统一一部分是 eV一部分是 Hartree或者混在同一列里。GNN 的输出没有任何物理约束它会自动学会拟合“当前这一批数字”但换到另一批单位就全乱。解决训练前对目标值 y 做数值范围检查打印 min、max、std。Hartree 转 eV 乘以 27.2114kcal/mol 转 eV 除以 23.06。我一般把单位换算单独写一个函数并在 Dataset 加载时断言 y 的范围在合理区间超了就报错。坑 2同样的模型加氢和不加氢结果完全不一样现象训练时用 AddHs推理时忘了加氢模型输出严重偏离甚至特征维度对不上。原因氢原子作为节点参与消息传递图的节点数和边数都变了模型学到的分布完全被破坏。解决把 add_hs 作为 Dataset 和推理脚本的固定参数两者必须来自同一个入口。做一致性检查同一个 SMILES 在训练前和推理前的节点数必须一致否则直接抛异常。这个检查 10 秒钟能省掉一整晚的排错。坑 3随机划分让测试集 MAE 变得虚低现象交叉验证 MAE 0.06 eV投到一批新骨架分子上误差直接翻三倍。原因随机划分把同骨架的类似分子同时分到训练集和测试集模型是在做结构联想不是泛化。解决按 Murcko 骨架分桶划分或至少按 CAS 号去重后随机划分。报告精度时注明划分方式否则不同论文之间没法横向比较。5.2 模型侧的两个坑过平滑与坐标混用坑 4消息传递层加到 5 层精度反而下降现象层数从 2 加到 5训练 loss 能下降验证 MAE 先降后升所有节点的表示越来越相似。原因分子图直径小苯环体系 2~3 跳就能覆盖全局。层数继续加深等于重复做同样的聚合节点特征被过度平均这就是典型的过平滑现象。解决分子能量预测一般 2~4 层就够全局信息靠读出层而不是靠加深网络。真需要更大感受野用跳跃连接或 JK-net 在层间加权组合。坑 5把 mol_2D.mol 当 3D 构象用预测构象能量一塌糊涂现象想用这套 GNN 预测某个分子的力场单点能输入 2D 结构文件结果 MAE 高达几个 kcal/mol完全不可用。原因.mol 文件里是 2D 坐标没有键长键角信息。GNN 学的是拓扑到能量的映射不是三维构象到能量的映射本质上有适用边界。解决先明确预测目标。QM9 类数据给出的是分子图对应的电子能量图结构就够如果目标是构象能量或势能面需要用 ETKDG 或 RDKit 生成 3D 构象再考虑 3D GNN 或等变网络。拿到新数据时我通常会按这个顺序过一遍打印 y 的 min/max/std随机抽三个 SMILES 检查加氢前后节点数统计骨架桶数量决定划分方式固定 hidden_dim128、num_layers3 快速跑通最后用单样本推理核对输出范围。这套习惯能避免 80% 的无效训练。6. 进阶用多任务改一版一次预测四个能量再做误差归因前面几章把单能量预测跑通了但实际项目里往往不只有一个目标。QM9 同时提供 U0、U298、H298、G298、HOMO、LUMO 等多项性质你当然可以训练六个模型但数据复用效率太低。更合理的做法是在最后一个隐藏层后加一个多任务回归头共享底层消息传递网络class MultiTaskEnergy(torch.nn.Module): def __init__(self, hidden_dim128, n_tasks4): super().__init__() # 复用前面 GNNMolEnergy只把它当成编码器 self.encoder GNNMolEnergy(in_dim4, hidden_dimhidden_dim, out_dimhidden_dim, num_layers3) self.heads torch.nn.ModuleList([ torch.nn.Linear(hidden_dim, 1) for _ in range(n_tasks) ]) def forward(self, x, edge_index, batch): h self.encoder(x, edge_index, batch) # 共享分子表示 return torch.cat([head(h) for head in self.heads], dim-1)关键点是每个任务只有一个线性头消息传递部分是共用的。U0 和 H298 高度相关共享表示能让数据互相补充小数据集下比六个独立模型稳得多。训练时把 L1Loss 改成对所有任务输出取平均即可。模型训完之后下一步不是急着调参而是做误差归因。把验证集里误差最大的 20 个分子挑出来按 Murcko 骨架聚类看模型到底对哪类结构无能为力。我见过最典型的场景是误差集中在含氮杂环体系上说明训练数据里这类结构太少需要回数据层面补样本而不是继续调学习率。这种黑匣子分析比看 loss 曲线有用得多。从那以后我每次拿到新的分子能量数据都强制走一遍完整流程确认单位固定加氢逻辑按骨架划分用一个 3 层 GIN 快速跑通再做多任务和误差归因。这套流程帮我省掉了大量试错时间也希望帮到你。本文还有配套的精品资源点击获取
