SchNetPack 2.0: A neural network toolbox for atomistic machine learning

TL;DR

SchNetPack 2.0 is an enhanced neural network toolkit for atomistic machine learning, featuring improved data handling, equivariant neural network modules, and a PyTorch-based molecular dynamics implementation, enabling flexible and complex molecular modeling tasks.

Contribution

The paper introduces SchNetPack 2.0 with new modules, an improved data pipeline, and integration with PyTorch Lightning and Hydra for flexible, extendable atomistic machine learning workflows.

Findings

Enhanced data pipeline improves training efficiency

Inclusion of equivariant neural network modules

Supports complex molecular dynamics simulations

Abstract

SchNetPack is a versatile neural networks toolbox that addresses both the requirements of method development and application of atomistic machine learning. Version 2.0 comes with an improved data pipeline, modules for equivariant neural networks as well as a PyTorch implementation of molecular dynamics. An optional integration with PyTorch Lightning and the Hydra configuration framework powers a flexible command-line interface. This makes SchNetPack 2.0 easily extendable with custom code and ready for complex training task such as generation of 3d molecular structures.