Logo
Project Icon

schnetpack

原子级系统深度学习建模工具包

SchNetPack深度神经网络原子系统量子化学分子动力学Github开源项目

SchNetPack是一个开源的深度学习工具包,用于原子级系统建模。它提供了构建和训练神经网络的基础组件,可预测分子和材料的势能面及量子化学性质。该工具包支持SchNet和PaiNN等先进模型,能够计算偶极矩、极化率等多种属性,并集成了分子动力学模拟功能。SchNetPack简化了新模型的开发和评估流程,为原子级机器学习研究提供了有力支持。

访问官网访问官网GithubGithub文档文档论文论文
介绍相关项目

SchNetPack - 原子系统深度神经网络

SchNetPack是一个用于开发和应用深度神经网络来预测分子和材料的势能面及其他量子化学性质的工具箱。它包含原子神经网络的基本构建模块，管理它们的训练，并提供对常见基准数据集的简单访问。这使得新模型的实现和评估变得容易。

文档可以在这里找到。

特性
  • SchNet - 用于分子和材料的端到端连续滤波CNN [1-3]
  • PaiNN - 用于分子和材料的等变消息传递 [4]
  • 用于偶极矩、极化率、应力和一般响应属性的输出模块
  • 静电、Ewald求和、ZBL排斥的模块
  • GPU加速的分子动力学代码，包括路径积分MD、恒温器、恒压器

安装

通过pip安装

安装SchNetPack最简单的方法是通过pip，它会自动从PyPI获取源代码：

pip install schnetpack

从源代码安装

你也可以从我们的仓库安装最新的代码：

git clone https://github.com/atomistic-machine-learning/schnetpack.git
cd schnetpack
pip install .

使用Tensorboard进行可视化

SchNetPack通过PyTorch Lightning支持多种日志后端。默认的日志器是Tensorboard。SchNetPack还支持TensorboardX。

入门

开始使用的最佳方式是通过命令行界面（CLI）在常见的基准数据集上训练SchNetPack模型。 安装SchNetPack时，训练脚本spktrain会被添加到你的PATH中。 CLI使用Hydra，并基于可以在这里找到的PyTorch Lightning/Hydra模板。 这使得模型、数据和训练过程的配置变得灵活。 为了充分利用这些功能，查看Hydra和PyTorch Lightning的文档可能会有帮助。

示例1：QM9

接下来，我们专注于使用CLI在QM9数据集上进行训练，但相同的程序也适用于其他基准数据集。 首先，创建一个工作目录，用于存储所有数据和运行：

mkdir spk_workdir
cd spk_workdir

然后，可以通过以下命令开始使用默认设置训练QM9的SchNet模型：

spktrain experiment=qm9_atomwise

脚本会打印实验配置qm9_atomwise的默认值。 如果数据集还不存在，它将自动下载到spk_workdir/data。 然后，训练将开始。

配置中的所有值都可以从命令行更改，包括运行和数据的目录。 默认情况下，模型存储在spk_workdir/runs的子目录中，该子目录具有唯一的运行ID哈希。 这可以按如下方式更改：

spktrain experiment=qm9_atomwise run.data_dir=/my/data/dir run.path=~/all_my_runs run.id=this_run

如果你调用spktrain experiment=qm9_atomwise --help，你可以看到完整的配置，其中包含所有可以更改的参数。 嵌套参数可以按如下方式更改：

spktrain experiment=qm9_atomwise run.data_dir=<path> data.batch_size=64

Hydra将参数组织在配置组中，这允许由多个yaml文件组成的分层配置。这使得轻松更改整个数据集、模型或表示变得可能。 例如，要从默认的SchNet表示更改为PaiNN，使用：

spktrain experiment=qm9_atomwise run.data_dir=<path> model/representation=painn

刚开始时，何时使用"."或"/"可能有些混淆。如果你正在加载预配置的配置组，则使用斜杠，而更改单个值时使用点。例如，配置组"model/representation"对应于配置的以下部分：

model:
  representation:
    _target_: schnetpack.representation.PaiNN
    n_atom_basis: 128
    n_interactions: 3
    shared_interactions: false
    shared_filters: false
    radial_basis:
      _target_: schnetpack.nn.radial.GaussianRBF
      n_rbf: 20
      cutoff: ${globals.cutoff}
    cutoff_fn:
      _target_: schnetpack.nn.cutoff.CosineCutoff
      cutoff: ${globals.cutoff}

如果你想额外更改该组的某些值,可以使用:

spktrain experiment=qm9_atomwise run.data_dir=<路径> model/representation=painn model.representation.n_interactions=5

有关配置组的更多详细信息,请查看Hydra文档。

示例2:势能面

上面的示例在内部使用AtomisticModel来预测单个属性,它是一个pytorch_lightning.LightningModule。以下示例将使用相同的类来预测势能面,特别是能量及其适当的导数以获得力和应力张量。这是可行的,因为通过命令行提供的MD17数据集的预定义配置experiment=md17选择了AtomisticModel使用的表示和输出模块。有关配置以及如何构建自定义配置的更详细描述可以在此处找到。

spktrain脚本可用于训练MD17数据集中某个分子的模型

spktrain experiment=md17 data.molecule=uracil

对于MD17,能量和力的参考计算是可用的。因此,需要为这些属性的损失设置权重。损失被定义为task配置组中输出定义的一部分:

    task:
      outputs:
        - _target_: schnetpack.task.ModelOutput
          name: ${globals.energy_key}
          loss_fn:
            _target_: torch.nn.MSELoss
          metrics:
            mae:
              _target_: torchmetrics.regression.MeanAbsoluteError
            mse:
              _target_: torchmetrics.regression.MeanSquaredError
          loss_weight: 0.005
        - _target_: schnetpack.task.ModelOutput
          name: ${globals.forces_key}
          loss_fn:
            _target_: torch.nn.MSELoss
          metrics:
            mae:
              _target_: torchmetrics.regression.MeanAbsoluteError
            mse:
              _target_: torchmetrics.regression.MeanSquaredError
          loss_weight: 0.995

对于能量和力的训练,我们建议在训练期间对力预测的损失赋予更大的权重。默认情况下,能量的损失权重设置为0.005,力的损失权重设置为0.995。这可以按如下方式更改:

spktrain experiment=md17 data.molecule=uracil task.outputs.0.loss_weight=0.005 task.outputs.1.loss_weight=0.995

日志记录

除了命令行输出外,SchNetPack还通过PyTorch Lightning支持多个日志后端。默认情况下,Tensorboard日志记录器处于激活状态。如果安装了TensorBoard,可以通过以下方式显示结果:

tensorboard --logdir=<运行目录>

此外,SchNetPack还提供了CSV日志记录器和Aim的配置。可以按如下方式选择:

spktrain experiment=md17 logger=csv

LAMMPS接口

SchNetPack提供了LAMMPS接口。详细的安装指南链接在我们文档的操作指南部分。

扩展

SchNetPack可以作为实现高级原子神经网络和训练任务的基础。例如,有一个名为schnetpack-gschnet的扩展包,用于最新版本的cG-SchNet [5],这是一个条件生成分子模型。它展示了如何在几个自定义类中实现复杂的训练任务,同时利用SchNetPack框架的分层配置和自动化训练程序。

引用

如果您在研究中使用SchNetPack,请引用:

K.T. Schütt, S.S.P. Hessmann, N.W.A. Gebauer, J. Lederer, M. Gastegger. SchNetPack 2.0: A neural network toolbox for atomistic machine learning. J. Chem. Phys. 2023, 158 (14): 144801. 10.1063/5.0138367。 K.T. Schütt, P. Kessel, M. Gastegger, K. Nicoli, A. Tkatchenko, K.-R. Müller. SchNetPack：原子系统深度学习工具箱。 J. Chem. Theory Comput. 2019, 15 (1): 448-455. 10.1021/acs.jctc.8b00908。

@article{schutt2023schnetpack,
    author = {Sch{\"u}tt, Kristof T. and Hessmann, Stefaan S. P. and Gebauer, Niklas W. A. and Lederer, Jonas and Gastegger, Michael},
    title = "{SchNetPack 2.0：原子机器学习的神经网络工具箱}",
    journal = {The Journal of Chemical Physics},
    volume = {158},
    number = {14},
    pages = {144801},
    year = {2023},
    month = {04},
    issn = {0021-9606},
    doi = {10.1063/5.0138367},
    url = {https://doi.org/10.1063/5.0138367},
    eprint = {https://pubs.aip.org/aip/jcp/article-pdf/doi/10.1063/5.0138367/16825487/144801\_1\_5.0138367.pdf},
}
@article{schutt2019schnetpack,
    author = {Sch{\"u}tt, Kristof T. and Kessel, Pan and Gastegger, Michael and Nicoli, Kim A. and Tkatchenko, Alexandre and Müller, Klaus-Robert},
    title = "{SchNetPack：原子系统深度学习工具箱}",
    journal = {Journal of Chemical Theory and Computation},
    volume = {15},
    number = {1},
    pages = {448-455},
    year = {2019},
    doi = {10.1021/acs.jctc.8b00908},
    URL = {https://doi.org/10.1021/acs.jctc.8b00908},
    eprint = {https://doi.org/10.1021/acs.jctc.8b00908},
}

致谢

PyTorch Lightning的CLI和hydra配置改编自此模板：

参考文献

  • [1] K.T. Schütt. F. Arbabzadah. S. Chmiela, K.-R. Müller, A. Tkatchenko. 深度张量神经网络的量子化学洞察。 Nature Communications 8. 13890 (2017) 10.1038/ncomms13890

  • [2] K.T. Schütt. P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, K.-R. Müller. SchNet：用于建模量子相互作用的连续滤波卷积神经网络。 神经信息处理系统进展 30, 第992-1002页 (2017) 论文

  • [3] K.T. Schütt. P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, K.-R. Müller. SchNet - 针对分子和材料的深度学习架构。 The Journal of Chemical Physics 148(24), 241722 (2018) 10.1063/1.5019779

  • [4] K. T. Schütt, O. T. Unke, M. Gastegger 用于预测张量属性和分子光谱的等变消息传递。 国际机器学习会议（第9377-9388页）。PMLR, 论文。

  • [5] N. W. A. Gebauer, M. Gastegger, S. S. P. Hessmann, K.-R. Müller, K. T. Schütt 使用条件生成神经网络逆向设计3D分子结构。 Nature Communications 13. 973 (2022) 10.1038/s41467-022-28526-y

相关项目
项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手，通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能，支持100+编程语言，与主流编辑器无缝集成，显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台，能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐，Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案，通过智能技术简化视频制作流程。无论是企业宣传还是个人分享，有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持，能够阅读和理解用户上传的文件内容，解析网页信息，并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题，Kimi都能以友好、专业的方式提供帮助。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术，为商家提供一键生成商品图和营销文案的服务，显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台，让商品第一时间被种草。

Project Cover

吐司

探索Tensor.Art平台的独特AI模型，免费访问各种图像生成与AI训练工具，从Stable Diffusion等基础模型开始，轻松实现创新图像生成。体验前沿的AI技术，推动个人和企业的创新发展。

Project Cover

SubCat字幕猫

SubCat字幕猫APP是一款创新的视频播放器，它将改变您观看视频的方式！SubCat结合了先进的人工智能技术，为您提供即时视频字幕翻译，无论是本地视频还是网络流媒体，让您轻松享受各种语言的内容。

Project Cover

美间AI

美间AI创意设计平台，利用前沿AI技术，为设计师和营销人员提供一站式设计解决方案。从智能海报到3D效果图，再到文案生成，美间让创意设计更简单、更高效。

Project Cover

AIWritePaper论文写作

AIWritePaper论文写作是一站式AI论文写作辅助工具，简化了选题、文献检索至论文撰写的整个过程。通过简单设定，平台可快速生成高质量论文大纲和全文，配合图表、参考文献等一应俱全，同时提供开题报告和答辩PPT等增值服务，保障数据安全，有效提升写作效率和论文质量。

使用协议隐私政策广告服务
投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号