Train NEP potentials in PyTorch

TorchNEP is a pure PyTorch implementation of the NEP4 (neuroevolution potential) training framework. Trained models are written as GPUMD nep.txt files and run directly in GPUMD, in LAMMPS through NEP_CPU, and in ASE.

Get started View on GitHub

pip install torchnep

from torchnep import train_nep

train_nep("nep.in", "train.xyz", output_dir="output")

GPUMD-compatible

nep.txt files load directly into GPUMD for molecular dynamics.

Two-stage training

A force-focused first stage, then an energy-focused second stage.

Multi-GPU and multi-node

Data-parallel training with near-linear scaling across nodes.

Fast on NVIDIA and AMD

torch.compile with automatic backend selection, tuned on CUDA and ROCm.

Memory-friendly

The dataset stays in host memory; GPU memory scales with the batch, not the dataset.

Fine-tuning

Start from any nep.txt or checkpoint, optionally slimmed to the elements of the new data.

Active learning

Choose the MD frames worth computing with DFT, with one model: the MaxVol extrapolation grade.

Plots

Loss curves, parity plots and error breakdowns straight from the output files.

Single-GPU training speed (left) and multi-node parallel scaling (right).

Ready-to-use examples — the training inputs and trained models of the TorchNEP paper — are in TorchNEP_models. If TorchNEP helps your research, please cite the paper.