简介基于GNN进行分子能量预测的Python完整源码与数据包面向希望将图神经网络应用于化学分子性质预测的研究者、学生或开发者覆盖从分子结构建模、特征提取到能量预测的完整实现链路。资源包共27个文件压缩包约5.4MB核心为8个Python脚本另含7个CSV数据文件、3个PT模型权重文件、3个TXT说明文件、2个MOL分子结构文件、2个PNG可视化结果图、1个Markdown文档及1个XML反应文件类型涵盖代码、数据、模型与输出结果。这份资料目前已有509人学习浏览。源码中包括数据加载、分子图转换、GNN模型定义、训练、评估与结果分析等模块并附有分子结构文件和训练结果图便于对照理解节点消息传递与聚合更新的细节。借助这些脚本读者可以快速跑通分子能量预测的完整流程还能进一步迁移到QM9等更大规模数据集上开展实验研究。1. GNN分子能量预测拿到源码先看单位和数据划分拿到「基于GNN进行分子能量预测Python完整源码和数据包」这套东西时我第一眼看的不是模型代码而是数据包里的能量单位。做GNN分子能量预测的Python项目十个里有五个翻车在单位换算、数据划分和原子坐标预处理上而不是网络结构本身。这篇笔记把从数据包到训练完成的完整链路拆开讲清楚怎么把xyz坐标和能量标签拼成图数据、用什么网络结构、损失函数怎么设、哪些参数是玄学但能救命的以及六个高发踩坑记录。适合刚入门分子机器学习的同学也适合从传统描述符方法转过来的人。2. 构造图数据把xyz数据包转成GNN能吃的分子图2.1 为什么是GNN能量是分子图的排列不变函数分子能量预测这个问题传统做法是手工计算描述符键长、键角、二面角、各类分子指纹再用随机森林或XGBoost拟合。这条路最大的问题不是精度而是描述符本身不可导——你没法从能量对原子坐标求梯度得到力更没法反向传播去优化几何结构。GNN解决的正是这个痛点。分子天然是一张图原子是节点化学键或空间近邻是边。能量是一个标量性质它只取决于分子几何不取决于分子在空间中的朝向。GNN的消息传递如果只用键长、键级这类旋转不变特征网络本身就能做到旋转等变下的不变性不需要数据增强去硬学这个物理规律。交换两个原子的编号图的邻接结构不变池化层把节点特征压成全局特征时用的是置换不变的求和或者求平均所以模型输出也不会变。这就是为什么标题里的源码会选择GNN而不是一维CNN或Transformer。分子的原子编号、化学键类型、三维坐标三者必须同时进入模型图结构是最自然的信息载体。做GNN分子能量预测时特征拼装、边的构建、池化方式这三个部分决定了模型的上限后面网络多两层少两层反而是次要的。2.2 数据包里的xyz/csv解析和归一化流程数据包拿到手常见的组织方式是三类文件原子坐标文件xyz格式、能量标签表csv、数据划分索引txt或json。xyz文件的第一行是原子数第二行是注释行一般会把能量等性质写在注释里后面跟着每个原子的元素符号和三维坐标。解析这件事尽量自己写不要迷信某个特定库因为不同来源的数据包注释行格式千差万别。def parse_xyz(path): with open(path, r, encodingutf-8) as f: lines f.read().strip().splitlines() n_atoms int(lines[0].strip()) comment lines[1].strip() atom_lines lines[2:2 n_atoms] atomic_numbers, coords [], [] for line in atom_lines: parts line.split() atomic_numbers.append(atom_to_number(parts[0])) # C - 6, H - 1 coords.append([float(x) for x in parts[1:4]]) energy parse_comment_energy(comment) # 从注释里提出能量值 return atomic_numbers, coords, energy这段代码的逻辑前三步处理文件头原子数决定后面读多少行注释行的格式不定可能是“energy-40.5 eV”也可能是“U0...”甚至没有能量此时能量要从外部的csv里单独读。atom_to_number是一张元素符号到原子序数的映射表H是1、C是6、N是7、O是8这张表在后续特征编码里也要用。坐标统一存成浮点数列表。参数说明解析技能里最容易被忽略的是注释行缺失的情况。数据包如果给的是单独的csv能量标签就要把坐标和能量按文件名或索引对齐不要假设注释行里一定有能量。归一化方面坐标的单位通常是Å埃如果数据包给的是nm要乘以10否则后面模型学到的长度尺度全是错的。能量单位需要确认是Hartree还是eV这两个单位差27.2倍是第一个大坑后面第五章详细展开。2.3 特征编码原子序数、键级和局部几何怎么做成张量图数据拼好了接下来的核心问题是每个节点和每条边拿什么特征进入网络。节点特征常见做法是原子序数做one-hot嵌入加上该原子的局部环境特征。原子序数嵌入是必须的C、N、O、H、S这几种元素在有机分子里出现频率最高嵌入维度取16已经够用。如果做的是金属有机体系建议把过渡金属单独建模不要和主族元素共用一套嵌入。局部环境特征比如该原子的配位数、是否在芳香环内、部分电荷这些特征从xyz里算不出的话就直接省掉不硬凑。边特征常见做法是键长1维、键级单键/双键/三键/芳香键的one-hot、键的共轭标记。键长是核心GNN的能量预测能力很大程度上依赖边特征里有没有连续可导的几何量。用表格整理一下我一般会进模型的字段对象特征维度说明节点原子序数嵌入16从元素符号映射覆盖前四周期节点局部电荷可选1有数据就加没有就用0边键长1单位Å必须和训练集统一边键级one-hot4单键/双键/三键/芳香键边共轭标记1是否参与共轭体系边的构建有两种路线。一种是按化学键用RDKit从分子结构式里出键的连接关系这条路对孤立的xyz坐标无效因为xyz里没有键的信息需要额外输入SMILES。另一种是按截断半径建图凡是在一定距离内的原子对都连一条边配一个键长属性。对能量预测任务我一般用截断半径法半径取4.0到5.0 Å之间能把非键相互作用也纳入消息传递范围这对能量预测很重要因为分子间的范德华作用在能量里贡献不小。主作者在数据包里如果没有预生成edge_index建议自己按截断半径生成然后用化学键表做一次校验看是不是把所有该连的键连上了。3. 搭建分子能量预测模型可运行的MPNN骨架与参数3.1 骨架拆分嵌入层、消息传递层、全局池化读出一个能用于分子能量预测的GNN网络结构可以拆成三段节点嵌入层、消息传递层、全局池化读出层。节点嵌入层负责把原子序数等离散特征映射到连续向量空间。消息传递层是核心每个原子把自己的特征发给邻居邻居收到后结合边特征做聚合再把聚合结果更新回自身特征。这个过程堆叠三到四层每层的感受野就到第3到4近邻壳层。读出层把图中所有节点的特征压成一个全局向量再通过MLP回归出能量。注意这里的关键选择池化用求和而不是求平均。分子的总能量是所有原子贡献的加和一个含有20个原子的分子和一个含40个原子的分子前者总能量大体上是后者的一半量级。用平均池化会把分子大小信息抹掉模型学到的其实是每原子平均能量用它预测总能量还要再乘原子数容易在边界的氢原子数统计上出偏差。求和池化配合MLP可以直接输出总能量这也是CGCNN和SchNet等模型的一贯做法。3.2 最小可运行代码一个MPNN风格的EnergyPredictor下面是一份能直接放进训练脚本里的模型骨架基于PyTorch和PyTorch Geometric实现。代码风格是标准的MPNN消息函数拼接源节点、目标节点和边特征再过一个MLP。import torch import torch.nn as nn import torch.nn.functional as F from torch_geometric.nn import MessagePassing, global_add_pool class MPNNLayer(MessagePassing): def __init__(self, node_dim: int, edge_dim: int): super().__init__(aggrmean) self.mlp nn.Sequential( nn.Linear(node_dim * 2 edge_dim, node_dim), nn.ReLU(), nn.Linear(node_dim, node_dim), ) def forward(self, x, edge_index, edge_attr): return self.propagate(edge_index, xx, edge_attredge_attr) def message(self, x_j, x_i, edge_attr): return self.mlp(torch.cat([x_i, x_j, edge_attr], dim-1)) class EnergyPredictor(nn.Module): def __init__(self, node_dim: int, hidden_dim: int, edge_dim: int, num_layers: int 3): super().__init__() self.embedding nn.Linear(node_dim, hidden_dim) self.layers nn.ModuleList([ MPNNLayer(hidden_dim, edge_dim) for _ in range(num_layers) ]) self.readout nn.Sequential( nn.Linear(hidden_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1), ) def forward(self, x, edge_index, edge_attr, batch): h self.embedding(x) for layer in self.layers: h F.silu(layer(h, edge_index, edge_attr)) h global_add_pool(h, batch) return self.readout(h).squeeze(-1)代码逻辑说明MPNNLayer继承MessagePassingaggrmean表示邻居消息聚合方式取平均稳定性和对节点度的适应性都比求和好。message函数接收x_j邻居特征、x_i中心原子特征和edge_attr三部分拼起来过一个两层MLP这就是一次消息的完整内容。forward里propagate自动完成邻居消息的聚合。EnergyPredictor里self.embedding处理原始节点特征num_layers默认3层每层后接SiLU激活。读出部分在global_add_pool之后接MLP输出维度是1回归能量值。为什么aggrmean和global_add_pool并存消息聚合用平均是为了稳定节点更新而最终读出用求和是为了还原总能量的加和性质这两层设计目标不同不冲突。参数调整建议hidden_dim取64到128之间分子数据集一个点的样本量在几万到十几万hidden_dim超过256参数量涨幅明显但精度提升有限。num_layers在3到5层之间效果最好超过6层会出现过平滑——所有节点的特征趋同能量预测精度反而下滑。edge_dim取决于边特征维度对应2.3节表格里是6维。3.3 输出头设计回归标量能量不是分类score输出层看起来是一行代码但隐藏着任务本质。能量预测是回归任务输出头必须是一个不带激活的线性层。很多人刚接触GNN时会把分类任务的交叉熵头顺手搬过来或者给最后一层加一个Sigmoid这是错误的Sigmoid把输出硬压到0到1之间而分子能量的范围在几十到几百电子伏特之间根本表达不了。还有一点读出层的MLP不要加归一化层。BatchNorm会打乱输出标量的绝对尺度因为能量需要的是精确的回归值不是归一化后的相对值。如果需要正则把Dropout放在消息传递层的MLP里而不是读出层。更极端的做法是直接在最后一层加一个常数偏置用来拟合训练集能量均值这在QM9数据集上能把初始loss降到接近零属于实践经验里的一个快速收敛技巧但前提是能量标签做过平移。如果你用的是预训练源码里的模型先确认它的输出层是nn.Linear(hidden, 1)而不是带激活函数的版本。一个带ReLU的输出头会在能量为负时直接把梯度截断分子总能量普遍是负值后果就是模型只能学到能量绝对值的分布方向完全学反。4. 训练配方损失函数、数据划分和优化器参数设置4.1 损失用每原子能量MAE评估才不会被分子大小干扰训练的时候用F.l1_loss还是F.mse_loss以及除以总原子数还是不予理会这两步各有讲究。先看损失函数选型。L1损失对离群点不敏感MSE损失会放大大误差样本的梯度。如果数据包是DFT算出来的能量存在少量几何优化不收敛的离群样本用L1更稳。如果确定数据干净MSE收敛更快。我一般在训练前600个epoch用L1后面用SmoothL1切换到MSE的平滑版本切换时机以验证集MAE不再下降为准。再看要不要除以原子数。理想情况下总能量的绝对误差应该只占总能量的一小部分但不同分子大小差异大直接对总能量做L1损失大分子在平均指标里天然占了更大权重模型的注意力会偏向大分子。更合理的做法是用每原子能量误差做损失def train_epoch(model, loader, optimizer, device): model.train() total_loss 0.0 for data in loader: data data.to(device) pred model(data.x, data.edge_index, data.edge_attr, data.batch) loss torch.nn.functional.l1_loss(pred, data.y) # 每原子能量误差版本把总能量误差除以各自原子数 # per_atom_loss (torch.abs(pred - data.y) / data.num_atoms).mean() optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * data.num_graphs return total_loss / len(loader.dataset)逻辑说明这里data.num_atoms是构造数据集时记录的张量形状和pred对齐。L1版本计算的是总能量的平均绝对误差而注释掉的每原子版本会先按每个分子自己的原子数归一再求平均。两版差异在混合了大小分子的批次里尤其明显建议训练时用每原子版本验证指标也统一用每原子MAE单位是eV/atom或者meV/atom。评估时还要注意批量大小对指标的影响。MAE是除以样本数还是除以原子总数会得到完全不同的数字。报告的时候必须写清楚是总能量MAE单位eV还是每原子MAE单位meV/atom否则不同论文之间的结果没法公平对比。4.2 按分子骨架划分train/val/test别直接随机打散这是分子能量预测项目里最容易被忽略、影响却最大的一个环节。直接train_test_split(random_state42)会让验证集里出现大量和训练集含相同骨架、只是取代基不同的分子。模型学会了记忆官能团组合验证集指标漂亮一旦换上真正的新分子就崩。正确的做法是按分子骨架划分。常见做法是用RDKit把分子标准化后取Bemis-Murcko骨架对所有分子做骨架聚类再把同一个骨架下的分子全部放到同一份数据里。这样划分出的测试集模型面对的是没见过的骨架指标才代表真实泛化能力。from rdkit import Chem from rdkit.Chem.Scaffolds import MurckoScaffold def scaffold_split(smiles_list, frac_train0.8, frac_val0.1, frac_test0.1): scaffold_groups {} for idx, smi in enumerate(smiles_list): mol Chem.MolFromSmiles(smi) if mol is None: continue scaffold MurckoScaffold.MurckoScaffoldSmiles( molmol, includeChiralityFalse ) scaffold_groups.setdefault(scaffold, []).append(idx) scaffold_list sorted(scaffold_groups.items(), keylambda kv: len(kv[1]), reverseTrue) train_idx, val_idx, test_idx [], [], [] train_count 0 for scaffold, idxs in scaffold_list: frac_done train_count / len(smiles_list) if frac_done frac_train: train_idx.extend(idxs) train_count len(idxs) elif frac_done frac_train frac_val: val_idx.extend(idxs) else: test_idx.extend(idxs) return train_idx, val_idx, test_idx这段代码的逻辑先按骨架把分子索引分组再按骨架大小为序依次填充训练集这样尽量让大的骨架簇先进入训练集防止某个骨架单独霸占测试集。参数frac_train、frac_val、frac_test控制划分比例骨架划分通常不需要严格按0.8/0.1/0.1执行按骨架占的样本量动态填充也可以关键是保证三个集合的骨架集合没有交集。如果手里的数据包没有提供SMILES只有xyz坐标退而求其次可以用一个简化替代按分子里的重原子数量和碳原子数量分箱在每个箱内随机划分。这比完全随机好因为至少保证测试集的分子大小分布和训练集不太一样但它的强度弱于骨架划分只在没有结构式输入时当做次优解。4.3 学习率、批大小和早停一组不翻车的基础参数优化器这里存在大量玄学但有一组参数可以当做后悔药兜底。学习率是分子GNN里最容易翻车的参数没有之一。lr1e-2起步几乎必然发散尤其是用Adam配合L1损失时梯度尺度不稳定会让loss直接变NaN。我的经验值Adam优化器下hidden_dim128、batch_size32时峰值学习率设在5e-4到1e-3之间。配合warmup策略前10个epoch从1e-4线性升到峰值之后用余弦退火降到1e-5左右。批大小直接决定显卡显存占用。分子GNN的batch是按图批量不是按节点批量。一个批里如果放了大量含50个以上重原子的分子显存占用会突然翻倍这源于PyG把同一个batch的所有图拼接成一张大图邻接矩阵的非零元素是各图边的加和。我通常用最大原子数控制动态组batch让每个batch的总边数不超过一个阈值比如5000到10000条边显存占用就非常可控。早停要看验证集MAE而不是训练loss。训练loss会一路下降到一个极低值但验证集MAE在某个epoch之后会开始上涨这个拐点就是过拟合开始的时候。用下面这段代码做兜底best_mae float(inf) patience 50 wait 0 for epoch in range(1, max_epochs 1): train_loss train_epoch(model, train_loader, optimizer, device) val_mae evaluate(model, val_loader, device) if val_mae best_mae: best_mae val_mae wait 0 torch.save(model.state_dict(), best_model.pt) else: wait 1 if wait patience: print(fearly stop at epoch {epoch}, best val mae {best_mae:.4f}) break逻辑说明patience50表示验证集MAE连续50个epoch没创新就停best_model.pt保存的是历史最优权重而不是最后一个epoch的权重。注意这个参数组合里的学习率要配合余弦退火调度器先用warmup冲到峰值再逐步降下来Early stopping往往在退火的后半段触发正好卡在验证损失的最低点附近。没有调度器的话patience要放宽到80因为模型在恒定学习率下震荡epoch间的验证指标波动较大动不动就会误停。5. 避坑分子能量GNN项目里最高频的六种问题5.1 能量单位不统一导致loss彻底失控现象训练loss一开始就是几千甚至NaN模型输出和标签差了3个数量级怎么调学习率都没用。原因数据包里的能量有的是Hartree标量约几到几十有的是eV几十到几百有的是kcal/mol几百到上千。三个量纲混在一个批次里L1损失的梯度尺度完全乱了。解决进模型之前统一转成eV。1 Hartree等于27.2114 eV1 Hartree等于627.5 kcal/mol1 eV约等于23.06 kcal/mol。我在预处理脚本里加一道校验打印训练集中能量的mean和std如果std大于50大概率还有Hartree和eV混存马上排查。5.2 XYZ文件里的氢原子是否齐全直接影响能量绝对值现象同一个分子两个数据包里能量相差十几eV但图结构一模一样。原因xyz文件里有些是显式氢含有所有H原子坐标有些经过处理把氢原子并到了重原子上。能量标签对应的是哪一种氢处理规则图数据就必须用哪一种混用就会错位。解决确认数据包注释里氢原子的计数规则统一用RDKit做一次加氢对齐。对有机分子显式氢是必须的而许多从头算数据集输出的是全原子坐标那就保持原样。关键是训练集和测试集必须来自同一个处理规则不要混着用。5.3 随机划分数据集验证集指标虚高现象模型在验证集MAE只有5 meV/atom一拿到真实新分子测试集就变成80 meV/atom。原因验证集里有大量train set的同骨架分子模型实际上在记忆原子组合而不是学习物理规律。解决改用骨架划分见4.2节的代码。没有SMILES时至少按分子量分箱后划分。测试集里一定要保证骨架完全不重叠否则指标没有意义。5.4 截断半径建出来的图是碎块能量预测突变现象验证集里个别分子的预测能量和真实值差几百eV查坐标没有异常模型结构没有异常。原因截断半径设成了3.0 Å甚至更小导致部分原子的邻居列表为空这些孤立节点没有收到任何消息池化后对总能量的贡献完全取决于自身嵌入特征物理上完全不合理。解决把截断半径加到4.5 Å以上并在建图后检查每个图是否连通。连通性检查很简单从任意节点做BFS看访问节点数是否等于该图的节点数。数据包里如果有分子的参考键连表也可以拿来做校验图里应该包含所有化学键对应的边再加空间近邻边。5.5 GNN层数过多出现过平滑现象网络从3层加到7层训练loss下降验证MAE反而上升而且不同分子的输出能量越来越接近。原因消息传递本质是邻域平滑层数越多每个节点的特征越趋向于全局平均节点特异性信息被抹平这是经典的过平滑问题。解决把层数控制在3到4层并检查节点表征的秩。如果torch.svd(h)计算出的奇异值衰减特别快说明大部分节点特征已经趋同。另一个补救手段是在消息传递层之间插入残差连接h h F.silu(layer(h, edge_index, edge_attr))能显著延后过平滑的出现。5.6 数据包里有重复分子训练集和测试集重叠现象训练集loss降得很好但测试集作弊严重——打印测试集里预测误差为0的样本发现和训练集某条数据几乎一模一样的几何。原因数据包的原始来源可能做过多次任务拼接同一个分子在不同能量档位下出现了多次。切分时没有对去重做处理。解决切分前用分子骨架加坐标的哈希值做去重。先对每个xyz用坐标归一化去掉平动后算MD5再把完全重复的条目删除。这一步步做完数据总量会明显减少但每个划分里的样本是真正独立的指标才经得起推敲。6. 从能量到力验证模型真的学到了势能面只看能量MAE还不足以证明模型可靠。一个能拟合能量的模型未必能给出正确的力——能量对坐标的负梯度。对分子动力学和数据增强来说力才是刚需。做法很简单在模型前向后面接一个torch.autograd.grad对坐标求导。from torch.autograd import grad coord data.pos.requires_grad_(True) energy model(data.x, data.edge_index, data.edge_attr, data.batch) forces -grad(energy.sum(), coord, create_graphTrue)[0]这段代码的逻辑energy.sum()保证梯度能流回所有坐标create_graphTrue让求出的力还能再被求导这是后续训练能量-力联合损失的前提。验证时把预测力和DFT参考力做MAE对比如果力误差大于能量误差一个数量级以上说明模型只拟合了能量的平均值没有学到原子间相互作用的细节——常见原因就是消息传递层里缺少键长对输出的敏感性检查一下边特征有没有真的进到消息里。我现在的习惯是任何能量预测模型上线前先跑一张误差散点图横轴是原子的局部配位环境配位数或Bader电荷纵轴是每原子能量误差。GNN的误差如果和配位数强相关说明它没有完全学会不同化学环境下的能量贡献我会回到特征工程重做。这条路我在不同分子数据集上走了很多遍越来越觉得GNN分子能量预测的精髓不在网络多深而在数据的物理一致性、划分的严密性和损失函数的尺度把握上。希望帮到你。本文还有配套的精品资源点击获取