TorchMD-Net

Overview

TorchMD-Net is a toolkit for training and using neural network potentials for molecular and biomolecular systems, rather than a single pretrained predictive model. Its potentials are exposed as PyTorch modules, and the README describes integration with ACEMD, OpenMM and TorchMD for molecular dynamics workflows. TensorNet and TensorNet2 are listed as available architectures; the Equivariant Transformer, Transformer and Graph Neural Network architectures are marked deprecated.

Training can be configured through YAML files or command-line arguments, with command-line values taking precedence. For custom NumPy datasets, the documented inputs are atom types and atomic coordinates, paired with energy labels, force labels or both. Alternative dataset implementations return torch-geometric Data objects with z and pos, plus y, neg_dy or both. These interfaces support workflows that learn molecular energies and coordinate derivatives from labelled structures. The repository also documents extension points for custom prior models and new architectures.

For inference, TorchMD-Net provides checkpoint loading and an ASE calculator interface, including examples for pretrained AceFF models. The loader automatically remaps older TensorNet and TensorNet2 checkpoint layouts when it detects the documented legacy marker. Multi-node training is supported, but the README requires equal GPU counts across nodes and warns about heterogeneous GPU architectures and possible multi-GPU hangs. The source excerpts do not establish predictive accuracy for a particular chemical domain. The repository's MIT code licence should not be treated as evidence of terms for separately distributed model weights or datasets.

Key Features

• Implements TensorNet and TensorNet2 neural network potentials as PyTorch modules; lists ET, T and GN architectures as deprecated.
• Supports training configuration through YAML files and command-line arguments, with command-line overrides.
• Loads custom NumPy datasets containing atom types, coordinates and energy or force labels, and supports torch-geometric dataset implementations.
• Provides pretrained checkpoint loading through `load_model` and an ASE interface through `TMDNETCalculator`, with AceFF examples.
• Automatically detects and remaps documented older TensorNet and TensorNet2 checkpoint tensor layouts, with explicit override options.
• Documents custom prior models, architecture extension and multi-node training with GPU-count constraints.

Use Cases

• Intended evaluation: train a molecular potential on a labelled collection of structures and assess energy and force predictions on held-out configurations.
• Intended evaluation: load an AceFF checkpoint through the documented ASE calculator workflow and assess its suitability for the intended molecular systems.
• Intended evaluation: investigate use of trained potentials in an ACEMD, OpenMM or TorchMD simulation workflow, checking the relevant integration requirements.
• Intended evaluation: prototype a custom prior or architecture using the repository's extension guidance and assess it with an appropriate validation dataset.

How to Use

1. Read the documentation and repository README. Choose between training a potential and loading an existing checkpoint; note which architectures are deprecated.
2. Follow the README's pip or conda/mamba installation guidance. Select a PyTorch accelerator build using the linked PyTorch instructions, rather than assuming the example CUDA choice fits your environment.
3. For training, prepare atom types, coordinates and energy or force labels using the documented custom dataset interface. Consult the example configurations, remembering that command-line settings override YAML values.
4. For pretrained inference, consult the AceFF examples or checkpoint-loading guidance. Review the documented legacy-layout handling before overriding automatic detection.
5. As an intended evaluation, assess predictions on representative held-out structures before simulation use. If expanding to multi-node training, follow the README's equal-GPU-count requirement and review its hardware and communication limitations.