概览
GemNet 是一种几何消息传递神经网络,旨在预测分子的总能量以及作用于其原子的力。此仓库提供的是 PyTorch 参考实现,而非原始的 TensorFlow 2 实现;README 另行链接了后者。该项目将模型与其 NeurIPS 2021 论文以及一篇后续探讨分子动力学中图神经网络势的论文关联起来。
仓库支持从模型训练到分子预测和模拟示例的工作流程。用户通过 config.yaml 或 config_seml.yaml 配置工作。train.ipynb 演示模型训练,而 predict.ipynb 展示如何对通过 ASE 加载的分子进行预测。文档说明的输出是分子能量和原子受力;来源摘录没有说明完整的输入模式、支持的化学体系范围或训练数据要求。对于集群工作流程,train_seml.py 将训练与 Sacred 和 SEML 集成,ase_example.ipynb 则演示分子动力学用法。
另一个准备步骤涉及初始化时的激活缩放。用户可以选择所提供的 scaling_factors.json,也可以使用 fit_scaling.py 计算因子;README 指出,这些因子由不同的 GemNet 变体共享。摘录确认了这些工作流程组件,但没有提供数值精度结果、运行时间测量或已展示的模拟稳定性限值。因此,潜在用户应针对自己的分子和模拟条件评估其适用性,而不应将这些示例视为某项特定应用已得到验证。
主要功能
- 使用几何消息传递神经网络预测分子总能量和各个原子所受的力。
- 提供 PyTorch 参考实现;原始 TensorFlow 2 实现另行链接。
- 包含用于训练的 `train.ipynb`,以及用于预测通过 ASE 加载的分子的 `predict.ipynb`。
- 通过搭配 Sacred 和 SEML 使用的 `train_seml.py` 支持集群训练工作流程。
- 包含 `ase_example.ipynb`,用于演示在分子动力学模拟中使用 GemNet。
- 提供预先计算的激活缩放因子 `scaling_factors.json`,以及通过 `fit_scaling.py` 计算因子的流程。
使用场景
- 预期评估:使用基于 ASE 的预测笔记本和适当的参考结果,评估模型对代表性分子的能量和原子受力预测。
- 预期评估:调整训练笔记本和配置文件,以用于分子能量与受力学习研究。
- 预期评估:探索 GemNet 在 ASE 分子动力学工作流程中的作用,并检查其在预期模拟条件下的表现。
- 预期评估:评估 Sacred 和 SEML 训练工作流程是否适用于在集群上组织模型训练实验。
使用方式
- 阅读仓库 README,了解 PyTorch 工作流程,以及它与另行链接的 TensorFlow 2 实现之间的区别。采用前请查看代码许可证。
- 检查软件包设置和仓库的依赖要求。设置文件声明 Python
>=3.8;来源摘录未能确认经过测试的环境。 - 对于笔记本工作流程,调整
config.yaml;对于文档所述的集群训练路径,调整config_seml.yaml。查阅仓库中的train.ipynb或train_seml.py。 - 按照 README 的说明,选择提供的
scaling_factors.json,或者根据fit_scaling.py计算初始化缩放因子。 - 使用
predict.ipynb查看对通过 ASE 加载的分子的预测。如果目标是模拟,接着查看ase_example.ipynb。作为预期评估步骤,在依赖这些结果之前,应将输出与适当的参考数据进行比较。