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

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

文档可以在这里找到 这里。

特性

• SchNet - 一种端到端连续滤波卷积神经网络，适用于分子和材料 [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

然后，可以使用默认设置训练一个 SchNet 模型来处理 QM9 数据集：

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=<path> model/representation=painn model.representation.n_interactions=5

有关配置组的更多详细信息，请参阅 Hydra 文档。

示例 2：势能面

上面的示例在内部使用了 AtomisticModel，它是一个 pytorch_lightning.LightningModule，用于预测单个属性。下面的示例将使用相同的类来预测势能面，特别是包含适当导数的能量，以便获得力和应力张量。之所以可行，是因为通过命令行参数 experiment=md17 提供的 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=<rundir>

此外，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：用于原子尺度机器学习的神经网络工具箱。 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. 基于深度张量神经网络的量子化学见解。 《自然·通讯》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——一种用于分子和材料的深度学习架构。 《化学物理学报》第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 利用条件生成神经网络进行三维分子结构的逆向设计。 《自然·通讯》第13卷，973页（2022年）10.1038/s41467-022-28526-y

SchNetPack 快速上手指南

SchNetPack 是一个用于开发和应用深度神经网络的工具箱，旨在预测分子和材料的势能面及其他量子化学性质。它内置了 SchNet、PaiNN 等主流原子神经网络模块，并提供了便捷的训练流程和基准数据集访问接口。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux 或 macOS (Windows 用户建议使用 WSL2)
• Python: 版本 3.8 或更高
• 依赖项:
  • PyTorch (需预先安装与您的 CUDA 版本匹配的 PyTorch)
  • PyTorch Lightning
  • Hydra (用于配置管理)

提示：建议先根据 PyTorch 官网 指引安装好 PyTorch，再安装 SchNetPack。

安装步骤

您可以通过 pip 直接安装稳定版，或者从源码安装最新版本。

方式一：使用 pip 安装（推荐）

这是最简单的安装方式，将自动从 PyPI 获取源代码。

pip install schnetpack

国内加速：如果您在中国大陆，建议使用清华或阿里镜像源以加快下载速度：

pip install schnetpack -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装

如果您需要最新的开发版功能：

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

基本使用

安装完成后，spktrain 命令行工具将自动添加到您的环境变量中。SchNetPack 基于 Hydra 配置系统，允许通过命令行灵活调整模型、数据和训练参数。

1. 快速训练示例 (QM9 数据集)

以下命令将使用默认配置在 QM9 数据集上训练一个 SchNet 模型。如果本地没有数据，脚本会自动下载。

# 创建工作目录
mkdir spk_workdir
cd spk_workdir

# 开始训练
spktrain experiment=qm9_atomwise

2. 自定义配置

您可以直接在命令行修改配置参数，无需编辑文件。

• 修改数据路径和运行目录：

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

• 修改超参数（如批次大小）：

  spktrain experiment=qm9_atomwise data.batch_size=64
  

• 切换模型架构（例如从 SchNet 切换到 PaiNN）： 注意：加载预配置组使用 /，修改具体参数值使用 .。

  spktrain experiment=qm9_atomwise model/representation=painn
  

• 组合修改（切换模型并调整层数）：

  spktrain experiment=qm9_atomwise model/representation=painn model.representation.n_interactions=5
  

3. 查看帮助与完整配置

若要查看所有可配置的参数列表：

spktrain experiment=qm9_atomwise --help

4. 查看训练日志

SchNetPack 默认使用 TensorBoard 记录训练过程。训练开始后，在新终端运行以下命令即可可视化结果：

tensorboard --logdir=<rundir>

(请将 <rundir> 替换为实际的运行目录路径，通常位于 spk_workdir/runs/ 下)

---

更多高级用法（如 MD17 势能面训练、LAMMPS 接口等）请参阅官方文档：https://schnetpack.readthedocs.io