简介：基于图神经网络的分子能量预测是深度学习在化学领域的重要应用，这份资源面向机器学习与化学交叉方向的研究者和学习者，提供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 和 .xml |
| processed/ | 预处理后的图缓存，PyG Dataset 自动读写 |
| smiles_CAS.txt / cas20210902.txt | SMILES 与 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, dtype=torch.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, dtype=torch.long).t().contiguous() edge_attr = torch.tensor(edge_attr, dtype=torch.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_hs=True。显式加氢后，乙醇 C2H6O 的图从 3 个重原子变成 9 个节点，模型输入维度没变，但图结构完全不同。训练时加了氢、推理时没加，模型精度会立刻崩掉，甚至维度都对不上。节点特征里 GetTotalNumHs 和 IsInRing 是 RDKit 中区分相似原子环境最便宜的组合，比如苯环上的碳和脂肪链上的碳，原子序数相同，但环标志和氢数能把它们分开。后续想加更丰富的特征，可以把原子序数改成可学习的 embedding，但先跑通流程不需要。
2.3 把数据集组织成 PyG 的 Dataset，processed 缓存到底存了什么
单分子转成图之后，下一步是把整个数据集组织成 PyG 能直接喂给 DataLoader 的对象。标准做法是继承 torch_geometric.data.Dataset：
from 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_hs=True) data = Data(x=graph["x"], edge_index=graph["edge_index"], edge_attr=graph["edge_attr"], y=torch.tensor([e], dtype=torch.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_dim=4, hidden_dim=128, out_dim=1, num_layers=3): 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_dim；global_add_pool 把同一个 batch 里所有节点按 batch 向量分组求和，输出 shape 是 (batch_size, hidden_dim)；最后的 regressor 把图级向量压到 1 维能量。整个模型输入是 DataLoader 吐出的 batch，batch 向量由 PyG 自动生成。
参数上，hidden_dim=128 是通用起点，小数据集可以降到 64。num_layers=3 是分子图的常规选择，不要一味加深，后面避坑章会专门讲。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_dim=4, hidden_dim=128, out_dim=1, num_layers=3) ckpt = torch.load("model_epoch100.pt", map_location="cpu") 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_hs=True) # 乙醇 data = Data(**graph) # 单分子也要构造 batch 向量，全 0 即可 batch = torch.zeros(data.x.size(0), dtype=torch.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 小。
单位问题比指标选择更致命。不同来源的数据能量单位可能完全不同：
| 单位 | 换算关系 | 常见场景 |
|---|---|---|
| eV | 1 eV = 23.06 kcal/mol | QM9 的能量、HOMO/LUMO |
| Hartree | 1 Hartree = 27.2114 eV | Gaussian 等量子化学软件输出 |
| 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_dim=4, hidden_dim=128, out_dim=1) optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=10) loss_fn = torch.nn.L1Loss() train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False) 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(f"epoch {epoch:3d} train_loss={total / len(train_dataset):.4f} val_mae={val_mae:.4f}")batch_size=64 对分子图来说算紧凑，分子平均 20~40 个节点，显存占用不大。如果数据集里混着大分子，要按原子数分桶，后面避坑章会讲。L1Loss 对异常值不敏感，符合能量预测的诉求。scheduler 用验证集 MAE 驱动，patience=10 表示 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_ratio=0.2): scaffolds = {} for idx, s in enumerate(smiles_list): mol = Chem.MolFromSmiles(s) if mol is None: continue # 取 Murcko 骨架作为桶 key，同一骨架的分子必须进同一份数据 scaf = MurckoScaffold.MurckoScaffoldSmiles(mol=mol) scaffolds.setdefault(scaf, []).append(idx) train_idx, val_idx = [], [] for scaf, idxs in sorted(scaffolds.items(), key=lambda kv: len(kv[1]), reverse=True): 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.2114，kcal/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_dim=128、num_layers=3 快速跑通，最后用单样本推理核对输出范围。这套习惯能避免 80% 的无效训练。
6. 进阶：用多任务改一版，一次预测四个能量，再做误差归因
前面几章把单能量预测跑通了，但实际项目里往往不只有一个目标。QM9 同时提供 U0、U298、H298、G298、HOMO、LUMO 等多项性质，你当然可以训练六个模型，但数据复用效率太低。更合理的做法是在最后一个隐藏层后加一个多任务回归头，共享底层消息传递网络：
class MultiTaskEnergy(torch.nn.Module): def __init__(self, hidden_dim=128, n_tasks=4): super().__init__() # 复用前面 GNNMolEnergy，只把它当成编码器 self.encoder = GNNMolEnergy(in_dim=4, hidden_dim=hidden_dim, out_dim=hidden_dim, num_layers=3) 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 快速跑通，再做多任务和误差归因。这套流程帮我省掉了大量试错时间，也希望帮到你。
本文还有配套的精品资源，点击获取
