TorchMD-NET

TorchMD-NET 提供最先进的神经网络势能（NNPs）以及训练它们的机制。它提供了几种 NNPs 的高效快速实现，并集成在 GPU 加速的分子动力学代码中，如 ACEMD、OpenMM 和 TorchMD。TorchMD-NET 将其 NNPs 作为 PyTorch 模块暴露出来。

文档

文档可在 https://torchmd-net.readthedocs.io 获取

可用架构

• 等变变换器 (ET)
• 变换器 (T)
• 图神经网络 (GN)
• TensorNet

安装

TorchMD-Net 可在 conda-forge 中获得，可以通过以下命令安装：

mamba install torchmd-net

我们推荐使用 Mamba 而不是 conda。

从源代码安装

TorchMD-Net 使用 pip 安装，但在此之前你需要安装一些依赖项。查看此文档页面。

使用

可以通过配置 yaml 文件或直接通过命令行参数指定训练参数。examples/ 中可以找到一些模型和数据集的架构和训练规范示例。注意，如果一个参数同时出现在 yaml 文件和命令行中，命令行版本优先。 可以通过设置 CUDA_VISIBLE_DEVICES 环境变量来选择 GPU。否则，可以使用参数 --ngpus 来选择要训练的 GPU 数量（-1，默认值，使用所有可用的 GPU 或 CUDA_VISIBLE_DEVICES 中指定的 GPU）。请记住，nvidia-smi 报告的 GPU ID 可能与 CUDA_VISIBLE_DEVICES 使用的不同。 例如，要在 QM9 数据集上训练等变变换器，使用论文中描述的架构和训练超参数，可以运行：

mkdir output
CUDA_VISIBLE_DEVICES=0 torchmd-train --conf torchmd-net/examples/ET-QM9.yaml --log-dir output/

运行 torchmd-train --help 查看所有可用选项及其描述。

预训练模型

有关如何加载预训练模型的说明，请参见此处。

创建新数据集

如果你想在自定义数据上进行训练，首先看看 torchmdnet.datasets.Custom，它提供了加载 NumPy 数据集的功能，包括原子类型和坐标，以及能量、力或两者作为标签。 或者，你可以按照 torch-geometric 实现数据集的方式实现自定义类。也就是说，派生 Dataset 或 InMemoryDataset 类并实现必要的函数（更多信息在此）。数据集必须返回 torch-geometric Data 对象，至少包含键 z（原子类型）和 pos（原子坐标），以及 y（标签）、neg_dy（标签相对于原子坐标的负导数）或两者。

自定义先验模型

除了实现自定义数据集类，还可以向模型添加自定义先验模型。这可以通过在 torchmdnet.priors 中实现新的先验模型类并添加参数 --prior-model <PriorModelName> 来完成。 作为示例，请看 torchmdnet.priors.Atomref。

多节点训练

为了在多个节点上训练模型，必须设置一些环境变量，这些变量提供了 PyTorch Lightning 所需的所有信息。以下我们提供了一个在两台各有两个 GPU 的机器上开始训练的示例 bash 脚本。脚本必须在每个节点上启动一次。一旦在所有节点上启动了 torchmd-train，将使用 NCCL 在节点之间建立网络连接。

除了环境变量外，还必须使用参数 --num-nodes 指定参与训练的节点数。

export NODE_RANK=0
export MASTER_ADDR=hostname1
export MASTER_PORT=12910

mkdir -p output
CUDA_VISIBLE_DEVICES=0,1 torchmd-train --conf torchmd-net/examples/ET-QM9.yaml.yaml --num-nodes 2 --log-dir output/

• NODE_RANK：表示节点索引的整数。主节点必须为 0，每个额外节点递增 1。
• MASTER_ADDR：主节点的主机名或 IP 地址。所有参与节点相同。
• MASTER_PORT：用于节点间通信的空闲网络端口。PyTorch Lightning 建议默认使用端口 12910。

已知限制

• 由于 PyTorch Lightning 计算所需 DDP 进程数的方式，所有节点必须使用相同数量的 GPU。否则，训练将无法开始或崩溃。
• 当混合使用不同 GPU 架构的节点时，我们观察到性能下降 50 倍（测试使用 RTX 2080 Ti 和 RTX 3090）。
• 某些 CUDA 系统在多 GPU 并行训练期间可能会挂起。尝试 export NCCL_P2P_DISABLE=1，这会禁用直接的 GPU 对等通信。

引用

如果你在研究中使用 TorchMD-NET，请引用以下论文：

主要参考文献

@misc{pelaez2024torchmdnet,
title={TorchMD-Net 2.0: Fast Neural Network Potentials for Molecular Simulations}, 
author={Raul P. Pelaez and Guillem Simeon and Raimondas Galvelis and Antonio Mirarchi and Peter Eastman and Stefan Doerr and Philipp Thölke and Thomas E. Markland and Gianni De Fabritiis},
year={2024},
eprint={2402.17660},
archivePrefix={arXiv},
primaryClass={cs.LG}
}

TensorNet

@inproceedings{simeon2023tensornet,
title={TensorNet: 用于高效学习分子势能的笛卡尔张量表示},
author={Guillem Simeon and Gianni De Fabritiis},
booktitle={第三十七届神经信息处理系统会议},
year={2023},
url={https://openreview.net/forum?id=BEHlPdBZ2e}
}

等变变换器

@inproceedings{
tholke2021equivariant,
title={用于基于神经网络的分子势能的等变变换器},
author={Philipp Th{\"o}lke and Gianni De Fabritiis},
booktitle={国际表示学习会议},
year={2022},
url={https://openreview.net/forum?id=zNHzqZ9wrRB}
}

图网络

@article{Majewski2023,
  title = {蛋白质热力学的机器学习粗粒度势能},
  volume = {14},
  ISSN = {2041-1723},
  url = {http://dx.doi.org/10.1038/s41467-023-41343-1},
  DOI = {10.1038/s41467-023-41343-1},
  number = {1},
  journal = {自然通讯},
  publisher = {施普林格科学与商业媒体有限责任公司},
  author = {Majewski, Maciej and Pérez, Adrià and Th\"{o}lke, Philipp and Doerr, Stefan and Charron, Nicholas E. and Giorgino, Toni and Husic, Brooke E. and Clementi, Cecilia and Noé, Frank and De Fabritiis, Gianni},
  year = {2023},
  month = sep 
}

开发者指南

实现新的架构

要实现新的架构，你需要遵循以下步骤： 1. 在torchmdnet.models中创建一个新的类，该类继承自torch.nn.Model。以TorchMD_ET为模板。这是模型的最小实现：

class MyModule(nn.Module):
  def __init__(self, parameter1, parameter2):
    super(MyModule, self).__init__()
    # 在此定义你的模型
    self.layer1 = nn.Linear(10, 10)
    ...
    # 在此初始化你的模型参数
    self.reset_parameters()

    def reset_parameters(self):
      # 在此初始化你的模型参数
      nn.init.xavier_uniform_(self.layer1.weight)
    ...
    
  def forward(self,
        z: Tensor, # 原子序数，形状 (n_atoms, 1)
        pos: Tensor, # 原子位置，形状 (n_atoms, 3)
        batch: Tensor, # 批次向量，形状 (n_atoms, 1)。同一分子中的所有原子具有相同的值且是连续的。
        q: Optional[Tensor] = None, # 原子电荷，形状 (n_atoms, 1)
        s: Optional[Tensor] = None, # 原子自旋，形状 (n_atoms, 1)
    ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
    # 在此定义你的前向传播
    scalar_features = ...
    vector_features = ...
    # 返回标量和向量特征，以及原子序数、位置和批次向量
    return scalar_features, vector_features, z, pos, batch

2. 将模型添加到torchmdnet.models.__init__.py中的__all__列表中。这将使测试能够检测到你的模型。 3. 通过添加新条目告诉models.model.create_model如何初始化你的模块，例如：

     elif args["model"] == "mymodule":
        from torchmdnet.models.torchmd_mymodule import MyModule
        is_equivariant = False # 如果你的模型是等变的，则设置为True
        representation_model = MyModule(
            parameter1=args["parameter1"],
            parameter2=args["parameter2"],
            **shared_args, # 通常所有模型共享的参数
        )

4. 将初始化你的模块所需的任何新参数添加到scripts.train.get_args中。例如：

   parser.add_argument('--parameter1', type=int, default=32, help='MyModule所需的参数1')
   ...

5. 在torchmd-net/examples中添加一个使用你的模型的示例配置文件。 6. 通过在tests.utils.load_example_args中添加一个案例来使测试使用你的配置文件。例如：

if model_name == "mymodule":
        config_file = join(dirname(dirname(__file__)), "examples", "MyModule-QM9.yaml")

此时，如果你的模块缺少某些功能，测试会告诉你，你可以添加它。如果你向包添加了新功能，请为其添加测试。

代码风格

我们使用black。在提交之前，请对你修改的文件运行black。

测试

要运行测试，请安装软件包并在存储库的根目录中运行pytest。测试是了解如何使用包的不同组件的良好知识来源。