August 16, 20250 citationsOpen Access

TorchANI 2.0: An extensible, high performance library for the design, training, and use of NN-IPs

IPIgnacio Pickering JXJiangeng Xue KHKate Huddleston

Key Points

TorchANI 2.0 significantly improves speed and performance in training and inference for ANI models, showing it's over faster than previous versions.
Benchmarks indicate that using TorchANI 2.0 allows for training physically constrained models which represent essential qualities of chemical systems effectively.
The new modular system in TorchANI 2.0 enables the addition of pairwise potentials, enhancing the flexibility of model design for various applications.
CUDA-accelerated optimization improves memory efficiency, facilitating better performance for network ensembles and large-scale calculations.

Abstract

In this work we introduce TorchANI 2.0, a significantly improved version of the free and open source TorchANI software package for training and evaluation of ANI (ANAKIN-ME) deep learning models. TorchANI 2.0 builds upon the foundation of its predecessor, while addressing its limitations and introducing new features. These changes greatly enhance its extensibility, performance, and suitability as a framework for developing models ready for molecular dynamics applications. These improvements include the introduction of a modular system to add arbitrary pairwise potentials to models, CUDA-accelerated optimization for faster and more memory-efficient calculation of local atomic features, and a batched system for better performance of network ensembles, among others. Our benchmarks demonstrate that TorchANI 2.0 achieves significant speedup over previous versions in both training and inference, and the library enhancements allow users to train physically constrained models that better represent important qualities of chemical systems. We demonstrate this by introducing three new ANI models that incorporate these features and evaluating their capabilities

Ask AI

Helpful

Bookmark

View Full Paper