UPF has been focused on the continuous development of TorchMD, our software for the development and usage of machine learning potentials for molecular simulations. TorchMD is divided into three different parts/applications, with their own code and repositories, all hosted in Github and freely available for everyone.
The first, TorchMD itself, which is an end-to-end differentiable molecular dynamics code. TorchMD-NET is a software for training and creating the neural network potentials. TorchMD-CG is a specific application to learn coarse-grained potentials, currently specialized in protein folding.
Technical specifications
At the moment, our efforts for application hardening are focused on TorchMD-NET, and trying to optimize both the speed of neural network training and evaluation in a production environment. Regarding the first problem, we have rewritten our current code and network architectures using PyTorch Geometric, a PyTorch-based library with optimized code to easily write, implement and run Graph Neural Networks (GNNs). Implementing TorchMD-NET with PyTorch geometric provided us with a ~3x speed up in training, which is especially significant if you are dealing with large training datasets.
Another issue we were encountering is the simulation speed of TorchMD with neural network potentials, particularly in our application for coarse-graining simulations. In order to gain some time, we have been using TorchScript to compile our models and switch from a pure Python code to a TorchScript program, optimizing the code and making it independent of Python, making it possible to run the scripts in a production environment where Python speed might be disadvantageous. By using TorchScript with our trained neural network potentials, we have given a ~1.8x speed-up in our coarse-grained simulations using TorchMD.
