23 May 2024 | Raul P. Pelaez, Guillem Simeon, Raimondas Galvelis, Antonio Mirarchi, Peter Eastman, Stefan Doerr, Philipp Thölke, Thomas E. Markland, Gianni De Fabritiis
TorchMD-Net 2.0 is a significant advancement in neural network potentials (NNPs) for molecular simulations, offering improved computational efficiency and versatility. The framework now supports TensorNet, an O(3)-equivariant model using rank-2 Cartesian tensor representations, and the Equivariant Transformer, which uses scalar and vector representations. These models achieve state-of-the-art accuracy on benchmark datasets. TorchMD-Net also integrates physical priors, enhancing its applicability in research. The framework is optimized for performance, with features like CUDA-optimized neighbor search, memory-efficient data loading, and compatibility with leading molecular dynamics packages such as OpenMM. TorchMD-Net is available via GitHub and includes a comprehensive documentation. The software supports both training and inference, with a modular design allowing for customization. Recent optimizations include CUDA graphs, which improve performance, especially for smaller workloads. The framework also supports periodic boundary conditions and efficient batch processing. TorchMD-Net's integration with OpenMM enables direct use of neural network potentials as force fields in molecular dynamics simulations. The framework is designed to be flexible, allowing researchers to develop new models beyond potential energy and forces. The software is available under a permissive MIT license and is accessible through conda-forge. The framework supports various datasets and provides tools for training and evaluating NNPs. The results show that TorchMD-Net achieves high accuracy and efficiency, with TensorNet demonstrating state-of-the-art performance on the QM9 dataset. The framework is also efficient in molecular simulations, with TensorNet 0L achieving a simulation speed of about 70 ns/day. Overall, TorchMD-Net represents a major step forward in the development of neural network potentials for molecular simulations, offering a versatile and efficient tool for researchers.TorchMD-Net 2.0 is a significant advancement in neural network potentials (NNPs) for molecular simulations, offering improved computational efficiency and versatility. The framework now supports TensorNet, an O(3)-equivariant model using rank-2 Cartesian tensor representations, and the Equivariant Transformer, which uses scalar and vector representations. These models achieve state-of-the-art accuracy on benchmark datasets. TorchMD-Net also integrates physical priors, enhancing its applicability in research. The framework is optimized for performance, with features like CUDA-optimized neighbor search, memory-efficient data loading, and compatibility with leading molecular dynamics packages such as OpenMM. TorchMD-Net is available via GitHub and includes a comprehensive documentation. The software supports both training and inference, with a modular design allowing for customization. Recent optimizations include CUDA graphs, which improve performance, especially for smaller workloads. The framework also supports periodic boundary conditions and efficient batch processing. TorchMD-Net's integration with OpenMM enables direct use of neural network potentials as force fields in molecular dynamics simulations. The framework is designed to be flexible, allowing researchers to develop new models beyond potential energy and forces. The software is available under a permissive MIT license and is accessible through conda-forge. The framework supports various datasets and provides tools for training and evaluating NNPs. The results show that TorchMD-Net achieves high accuracy and efficiency, with TensorNet demonstrating state-of-the-art performance on the QM9 dataset. The framework is also efficient in molecular simulations, with TensorNet 0L achieving a simulation speed of about 70 ns/day. Overall, TorchMD-Net represents a major step forward in the development of neural network potentials for molecular simulations, offering a versatile and efficient tool for researchers.