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