TorchMD-NET

TorchMD-NET is an open-source Python library for researchers who train neural network potentials for molecular dynamics. It exposes its models as PyTorch modules and connects them to GPU-accelerated simulation software, including ACEMD, OpenMM and TorchMD. It's available under the MIT license.

TensorNet and TensorNet2 are its supported architectures. Researchers can train potentials on their own molecular data or use pretrained AceFF models for inference. It also loads checkpoints from models trained with the library, so work can continue with an existing trained potential.

Custom training data can include atom types and coordinates, with energies, forces or both as targets. The library accepts NumPy datasets and supports custom datasets built around PyTorch Geometric. Researchers who need a different model design can extend the architecture or add a custom prior model rather than stay within the supplied implementations.

Training can span multiple machines and GPUs through PyTorch Lightning. Each machine must use the same number of GPUs, and mixing GPU architectures can reduce performance. This makes the library relevant to teams with their own GPU training infrastructure as well as researchers integrating learned potentials into a simulation workflow.

For existing AceFF and TensorNet checkpoints, the model loader automatically handles an older internal tensor layout.